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

2
src/primaite/config/_package_data/example_config.yaml
View File

@@ -693,7 +693,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
nics:
network_interfaces:
2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot
ip_address: 192.168.10.110
subnet_mask: 255.255.255.0

2
src/primaite/config/_package_data/example_config_2_rl_agents.yaml
View File

@@ -1073,7 +1073,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
nics:
network_interfaces:
2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot
ip_address: 192.168.10.110
subnet_mask: 255.255.255.0

8
src/primaite/game/agent/actions.py
View File

@@ -550,7 +550,7 @@ class NetworkNICAbstractAction(AbstractAction):
nic_num = self.manager.get_nic_num_by_idx(node_idx=node_id, nic_idx=nic_id)
if node_name is None or nic_num is None:
return ["do_nothing"]
return ["network", "node", node_name, "nic", nic_num, self.verb]
return ["network", "node", node_name, "network_interface", nic_num, self.verb]
class NetworkNICEnableAction(NetworkNICAbstractAction):
@@ -723,8 +723,8 @@ class ActionManager:
node_obj = self.game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
nics = node_obj.nics
for nic_uuid, nic_obj in nics.items():
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
self.ip_address_list.append(nic_obj.ip_address)
# action_args are settings which are applied to the action space as a whole.
@@ -964,7 +964,7 @@ class ActionManager:
node_name = entry["node_name"]
nic_num = entry["nic_num"]
node_obj = game.simulation.network.get_node_by_hostname(node_name)
ip_address = node_obj.ethernet_port[nic_num].ip_address
ip_address = node_obj.network_interface[nic_num].ip_address
ip_address_list.append(ip_address)
obj = cls(

34
src/primaite/game/agent/observations.py
View File

@@ -406,7 +406,7 @@ class NodeObservation(AbstractObservation):
where: Optional[Tuple[str]] = None,
services: List[ServiceObservation] = [],
folders: List[FolderObservation] = [],
nics: List[NicObservation] = [],
network_interfaces: List[NicObservation] = [],
logon_status: bool = False,
num_services_per_node: int = 2,
num_folders_per_node: int = 2,
@@ -429,9 +429,9 @@ class NodeObservation(AbstractObservation):
:type folders: Dict[int,str], optional
:param max_folders: Max number of folders in this node's obs space, defaults to 2
:type max_folders: int, optional
:param nics: Mapping between position in observation space and NIC idx, defaults to {}
:type nics: Dict[int,str], optional
:param max_nics: Max number of NICS in this node's obs space, defaults to 5
:param network_interfaces: Mapping between position in observation space and NIC idx, defaults to {}
:type network_interfaces: Dict[int,str], optional
:param max_nics: Max number of network interfaces in this node's obs space, defaults to 5
:type max_nics: int, optional
"""
super().__init__()
@@ -456,11 +456,11 @@ class NodeObservation(AbstractObservation):
msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}"
_LOGGER.warning(msg)
self.nics: List[NicObservation] = nics
while len(self.nics) < num_nics_per_node:
self.nics.append(NicObservation())
while len(self.nics) > num_nics_per_node:
truncated_nic = self.nics.pop()
self.network_interfaces: List[NicObservation] = network_interfaces
while len(self.network_interfaces) < num_nics_per_node:
self.network_interfaces.append(NicObservation())
while len(self.network_interfaces) > num_nics_per_node:
truncated_nic = self.network_interfaces.pop()
msg = f"Too many NICs in Node observation for node. Truncating service {truncated_nic.where[-1]}"
_LOGGER.warning(msg)
@@ -469,7 +469,7 @@ class NodeObservation(AbstractObservation):
self.default_observation: Dict = {
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
"FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)},
"NICS": {i + 1: n.default_observation for i, n in enumerate(self.nics)},
"NETWORK_INTERFACES": {i + 1: n.default_observation for i, n in enumerate(self.network_interfaces)},
"operating_status": 0,
}
if self.logon_status:
@@ -494,7 +494,9 @@ class NodeObservation(AbstractObservation):
obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
obs["operating_status"] = node_state["operating_state"]
obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)}
obs["NETWORK_INTERFACES"] = {
i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces)
}
if self.logon_status:
obs["logon_status"] = 0
@@ -508,7 +510,9 @@ class NodeObservation(AbstractObservation):
"SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
"FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
"operating_status": spaces.Discrete(5),
"NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}),
"NETWORK_INTERFACES": spaces.Dict(
{i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}
),
}
if self.logon_status:
space_shape["logon_status"] = spaces.Discrete(3)
@@ -564,13 +568,13 @@ class NodeObservation(AbstractObservation):
]
# create some configs for the NIC observation in the format {"nic_num":1}, {"nic_num":2}, {"nic_num":3}, etc.
nic_configs = [{"nic_num": i for i in range(num_nics_per_node)}]
nics = [NicObservation.from_config(config=c, game=game, parent_where=where) for c in nic_configs]
network_interfaces = [NicObservation.from_config(config=c, game=game, parent_where=where) for c in nic_configs]
logon_status = config.get("logon_status", False)
return cls(
where=where,
services=services,
folders=folders,
nics=nics,
network_interfaces=network_interfaces,
logon_status=logon_status,
num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node,
@@ -728,7 +732,7 @@ class AclObservation(AbstractObservation):
node_ref = ip_map_config["node_hostname"]
nic_num = ip_map_config["nic_num"]
node_obj = game.simulation.network.nodes[game.ref_map_nodes[node_ref]]
nic_obj = node_obj.ethernet_port[nic_num]
nic_obj = node_obj.network_interface[nic_num]
node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2
router_hostname = config["router_hostname"]

9
src/primaite/game/agent/rewards.py
View File

@@ -68,6 +68,8 @@ class DummyReward(AbstractReward):
:param config: dict of options for the reward component's constructor. Should be empty.
:type config: dict
:return: The reward component.
:rtype: DummyReward
"""
return cls()
@@ -230,7 +232,12 @@ class WebpageUnavailablePenalty(AbstractReward):
@classmethod
def from_config(cls, config: dict) -> AbstractReward:
"""Build the reward component object from config."""
"""
Build the reward component object from config.
:param config: Configuration dictionary.
:type config: Dict
"""
node_hostname = config.get("node_hostname")
return cls(node_hostname=node_hostname)

109
src/primaite/game/game.py
View File

@@ -11,14 +11,17 @@ from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAge
from primaite.game.agent.observations import ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.session.io import SessionIO, SessionIOSettings
from primaite.simulator.network.hardware.base import NIC, NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import Router
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.hardware.base import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
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.applications.red_applications.dos_bot import DoSBot
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
@@ -31,6 +34,24 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
_LOGGER = getLogger(__name__)
APPLICATION_TYPES_MAPPING = {
"WebBrowser": WebBrowser,
"DatabaseClient": DatabaseClient,
"DataManipulationBot": DataManipulationBot,
"DoSBot": DoSBot,
}
SERVICE_TYPES_MAPPING = {
"DNSClient": DNSClient,
"DNSServer": DNSServer,
"DatabaseService": DatabaseService,
"WebServer": WebServer,
"FTPClient": FTPClient,
"FTPServer": FTPServer,
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
class PrimaiteGameOptions(BaseModel):
"""
@@ -231,54 +252,48 @@ class PrimaiteGame:
new_service = None
service_ref = service_cfg["ref"]
service_type = service_cfg["type"]
service_types_mapping = {
"DNSClient": DNSClient, # key is equal to the 'name' attr of the service class itself.
"DNSServer": DNSServer,
"DatabaseClient": DatabaseClient,
"DatabaseService": DatabaseService,
"WebServer": WebServer,
"FTPClient": FTPClient,
"FTPServer": FTPServer,
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
if service_type in service_types_mapping:
if service_type in SERVICE_TYPES_MAPPING:
_LOGGER.debug(f"installing {service_type} on node {new_node.hostname}")
new_node.software_manager.install(service_types_mapping[service_type])
new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service.uuid
else:
_LOGGER.warning(f"service type not found {service_type}")
# service-dependent options
if service_type == "DatabaseClient":
if service_type == "DNSClient":
if "options" in service_cfg:
opt = service_cfg["options"]
if "db_server_ip" in opt:
new_service.configure(server_ip_address=IPv4Address(opt["db_server_ip"]))
if "dns_server" in opt:
new_service.dns_server = IPv4Address(opt["dns_server"])
if service_type == "DNSServer":
if "options" in service_cfg:
opt = service_cfg["options"]
if "domain_mapping" in opt:
for domain, ip in opt["domain_mapping"].items():
new_service.dns_register(domain, ip)
new_service.dns_register(domain, IPv4Address(ip))
if service_type == "DatabaseService":
if "options" in service_cfg:
opt = service_cfg["options"]
if "backup_server_ip" in opt:
new_service.configure_backup(backup_server=IPv4Address(opt["backup_server_ip"]))
new_service.configure_backup(backup_server=IPv4Address(opt.get("backup_server_ip")))
new_service.start()
if service_type == "FTPServer":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.server_password = opt.get("server_password")
new_service.start()
if service_type == "NTPClient":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.ntp_server = IPv4Address(opt.get("ntp_server_ip"))
new_service.start()
if "applications" in node_cfg:
for application_cfg in node_cfg["applications"]:
new_application = None
application_ref = application_cfg["ref"]
application_type = application_cfg["type"]
application_types_mapping = {
"WebBrowser": WebBrowser,
"DataManipulationBot": DataManipulationBot,
}
if application_type in application_types_mapping:
new_node.software_manager.install(application_types_mapping[application_type])
if application_type in APPLICATION_TYPES_MAPPING:
new_node.software_manager.install(APPLICATION_TYPES_MAPPING[application_type])
new_application = new_node.software_manager.software[application_type]
game.ref_map_applications[application_ref] = new_application.uuid
else:
@@ -294,12 +309,32 @@ class PrimaiteGame:
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
data_manipulation_p_of_success=float(opt.get("data_manipulation_p_of_success", "0.1")),
)
elif application_type == "DatabaseClient":
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.configure(
server_ip_address=IPv4Address(opt.get("db_server_ip")),
server_password=opt.get("server_password"),
)
elif application_type == "WebBrowser":
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.target_url = opt.get("target_url")
if "nics" in node_cfg:
for nic_num, nic_cfg in node_cfg["nics"].items():
elif application_type == "DoSBot":
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.configure(
target_ip_address=IPv4Address(opt.get("target_ip_address")),
target_port=Port(opt.get("target_port", Port.POSTGRES_SERVER.value)),
payload=opt.get("payload"),
repeat=bool(opt.get("repeat")),
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
dos_intensity=float(opt.get("dos_intensity", "1.0")),
max_sessions=int(opt.get("max_sessions", "1000")),
)
if "network_interfaces" in node_cfg:
for nic_num, nic_cfg in node_cfg["network_interfaces"].items():
new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"]))
net.add_node(new_node)
@@ -311,13 +346,13 @@ class PrimaiteGame:
node_a = net.nodes[game.ref_map_nodes[link_cfg["endpoint_a_ref"]]]
node_b = net.nodes[game.ref_map_nodes[link_cfg["endpoint_b_ref"]]]
if isinstance(node_a, Switch):
endpoint_a = node_a.switch_ports[link_cfg["endpoint_a_port"]]
endpoint_a = node_a.network_interface[link_cfg["endpoint_a_port"]]
else:
endpoint_a = node_a.ethernet_port[link_cfg["endpoint_a_port"]]
endpoint_a = node_a.network_interface[link_cfg["endpoint_a_port"]]
if isinstance(node_b, Switch):
endpoint_b = node_b.switch_ports[link_cfg["endpoint_b_port"]]
endpoint_b = node_b.network_interface[link_cfg["endpoint_b_port"]]
else:
endpoint_b = node_b.ethernet_port[link_cfg["endpoint_b_port"]]
endpoint_b = node_b.network_interface[link_cfg["endpoint_b_port"]]
new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)
game.ref_map_links[link_cfg["ref"]] = new_link.uuid

49
src/primaite/notebooks/uc2_demo.ipynb
View File

@@ -127,7 +127,7 @@
" - FILES\n",
" - <file_id 1-1>\n",
" - health_status\n",
" - NICS\n",
" - NETWORK_INTERFACES\n",
" - <nic_id 1-2>\n",
" - nic_status\n",
" - operating_status\n",
@@ -181,7 +181,7 @@
"\n",
"The ACL rules in the observation space appear in the same order that they do in the actual ACL. Though, only the first 10 rules are shown, there are default rules lower down that cannot be changed by the agent. The extra rules just allow the network to function normally, by allowing pings, ARP traffic, etc.\n",
"\n",
"Most nodes have only 1 nic, so the observation for those is placed at NIC index 1 in the observation space. Only the security suite has 2 NICs, the second NIC in the observation space is the one that connects the security suite with swtich_2.\n",
"Most nodes have only 1 network_interface, so the observation for those is placed at NIC index 1 in the observation space. Only the security suite has 2 NICs, the second NIC in the observation space is the one that connects the security suite with swtich_2.\n",
"\n",
"The meaning of the services' operating_state is:\n",
"|operating_state|label|\n",
@@ -314,7 +314,8 @@
"The blue agent's reward is calculated using two measures:\n",
"1. Whether the database file is in a good state (+1 for good, -1 for corrupted, 0 for any other state)\n",
"2. Whether each green agents' most recent webpage request was successful (+1 for a `200` return code, -1 for a `404` return code and 0 otherwise).\n",
"The file status reward and the two green-agent-related reward are averaged to get a total step reward.\n"
"\n",
"The file status reward and the two green-agent-related rewards are averaged to get a total step reward.\n"
]
},
{
@@ -334,7 +335,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
@@ -344,7 +347,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Imports\n",
@@ -367,7 +372,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# create the env\n",
@@ -415,7 +422,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"for step in range(35):\n",
@@ -433,7 +442,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"pprint(obs['NODES'])"
@@ -449,7 +460,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(9) # scan database file\n",
@@ -475,7 +488,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
@@ -500,7 +515,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
@@ -523,7 +540,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"env.step(13) # Patch the database\n",
@@ -573,7 +592,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -587,9 +606,9 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
"nbformat_minor": 4
}

4
src/primaite/simulator/__init__.py
View File

@@ -12,8 +12,8 @@ class _SimOutput:
self._path: Path = (
_PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
)
self.save_pcap_logs: bool = False
self.save_sys_logs: bool = False
self.save_pcap_logs: bool = True
self.save_sys_logs: bool = True
@property
def path(self) -> Path:

6
src/primaite/simulator/core.py
View File

@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
from typing import Callable, ClassVar, Dict, List, Optional, Union
from uuid import uuid4
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, Field
from primaite import getLogger
@@ -150,12 +150,10 @@ class SimComponent(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
"""Configure pydantic to allow arbitrary types and to let the instance have attributes not present in model."""
uuid: str
uuid: str = Field(default_factory=lambda: str(uuid4()))
"""The component UUID."""
def __init__(self, **kwargs):
if not kwargs.get("uuid"):
kwargs["uuid"] = str(uuid4())
super().__init__(**kwargs)
self._request_manager: RequestManager = self._init_request_manager()
self._parent: Optional["SimComponent"] = None

307
src/primaite/simulator/network/airspace.py Normal file
View File

@@ -0,0 +1,307 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Dict, Final, List, Optional
from prettytable import PrettyTable
from primaite import getLogger
from primaite.simulator.network.hardware.base import Layer3Interface, NetworkInterface, WiredNetworkInterface
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.core.packet_capture import PacketCapture
_LOGGER = getLogger(__name__)
__all__ = ["AIR_SPACE", "AirSpaceFrequency", "WirelessNetworkInterface", "IPWirelessNetworkInterface"]
class AirSpace:
"""Represents a wireless airspace, managing wireless network interfaces and handling wireless transmission."""
def __init__(self):
self._wireless_interfaces: Dict[str, WirelessNetworkInterface] = {}
self._wireless_interfaces_by_frequency: Dict[AirSpaceFrequency, List[WirelessNetworkInterface]] = {}
def show(self, frequency: Optional[AirSpaceFrequency] = None):
"""
Displays a summary of wireless interfaces in the airspace, optionally filtered by a specific frequency.
:param frequency: The frequency band to filter devices by. If None, devices for all frequencies are shown.
"""
table = PrettyTable()
table.field_names = ["Connected Node", "MAC Address", "IP Address", "Subnet Mask", "Frequency", "Status"]
# If a specific frequency is provided, filter by it; otherwise, use all frequencies.
frequencies_to_show = [frequency] if frequency else self._wireless_interfaces_by_frequency.keys()
for freq in frequencies_to_show:
interfaces = self._wireless_interfaces_by_frequency.get(freq, [])
for interface in interfaces:
status = "Enabled" if interface.enabled else "Disabled"
table.add_row(
[
interface._connected_node.hostname, # noqa
interface.mac_address,
interface.ip_address if hasattr(interface, "ip_address") else None,
interface.subnet_mask if hasattr(interface, "subnet_mask") else None,
str(freq),
status,
]
)
print(table)
def add_wireless_interface(self, wireless_interface: WirelessNetworkInterface):
"""
Adds a wireless network interface to the airspace if it's not already present.
:param wireless_interface: The wireless network interface to be added.
"""
if wireless_interface.mac_address not in self._wireless_interfaces:
self._wireless_interfaces[wireless_interface.mac_address] = wireless_interface
if wireless_interface.frequency not in self._wireless_interfaces_by_frequency:
self._wireless_interfaces_by_frequency[wireless_interface.frequency] = []
self._wireless_interfaces_by_frequency[wireless_interface.frequency].append(wireless_interface)
def remove_wireless_interface(self, wireless_interface: WirelessNetworkInterface):
"""
Removes a wireless network interface from the airspace if it's present.
:param wireless_interface: The wireless network interface to be removed.
"""
if wireless_interface.mac_address in self._wireless_interfaces:
self._wireless_interfaces.pop(wireless_interface.mac_address)
self._wireless_interfaces_by_frequency[wireless_interface.frequency].remove(wireless_interface)
def clear(self):
"""
Clears all wireless network interfaces and their frequency associations from the airspace.
After calling this method, the airspace will contain no wireless network interfaces, and transmissions cannot
occur until new interfaces are added again.
"""
self._wireless_interfaces.clear()
self._wireless_interfaces_by_frequency.clear()
def transmit(self, frame: Frame, sender_network_interface: WirelessNetworkInterface):
"""
Transmits a frame to all enabled wireless network interfaces on a specific frequency within the airspace.
This ensures that a wireless interface does not receive its own transmission.
:param frame: The frame to be transmitted.
:param sender_network_interface: The wireless network interface sending the frame. This interface will be
excluded from the list of receivers to prevent it from receiving its own transmission.
"""
for wireless_interface in self._wireless_interfaces_by_frequency.get(sender_network_interface.frequency, []):
if wireless_interface != sender_network_interface and wireless_interface.enabled:
wireless_interface.receive_frame(frame)
AIR_SPACE: Final[AirSpace] = AirSpace()
"""
A singleton instance of the AirSpace class, representing the global wireless airspace.
This instance acts as the central management point for all wireless communications within the simulated network
environment. By default, there is only one airspace in the simulation, making this variable a singleton that
manages the registration, removal, and transmission of wireless frames across all wireless network interfaces configured
in the simulation. It ensures that wireless frames are appropriately transmitted to and received by wireless
interfaces based on their operational status and frequency band.
"""
class AirSpaceFrequency(Enum):
"""Enumeration representing the operating frequencies for wireless communications."""
WIFI_2_4 = 2.4e9
"""WiFi 2.4 GHz. Known for its extensive range and ability to penetrate solid objects effectively."""
WIFI_5 = 5e9
"""WiFi 5 GHz. Known for its higher data transmission speeds and reduced interference from other devices."""
def __str__(self) -> str:
if self == AirSpaceFrequency.WIFI_2_4:
return "WiFi 2.4 GHz"
elif self == AirSpaceFrequency.WIFI_5:
return "WiFi 5 GHz"
else:
return "Unknown Frequency"
class WirelessNetworkInterface(NetworkInterface, ABC):
"""
Represents a wireless network interface in a network device.
This abstract base class models wireless network interfaces, encapsulating properties and behaviors specific to
wireless connectivity. It provides a framework for managing wireless connections, including signal strength,
security protocols, and other wireless-specific attributes and methods.
Wireless network interfaces differ from wired ones in their medium of communication, relying on radio frequencies
for data transmission and reception. This class serves as a base for more specific types of wireless network
interfaces, such as Wi-Fi adapters or radio network interfaces, ensuring that essential wireless functionality is
defined and standardised.
Inherits from:
- NetworkInterface: Provides basic network interface properties and methods.
As an abstract base class, it requires subclasses to implement specific methods related to wireless communication
and may define additional properties and methods specific to wireless technology.
"""
frequency: AirSpaceFrequency = AirSpaceFrequency.WIFI_2_4
def enable(self):
"""Attempt to enable the network interface."""
if self.enabled:
return
if not self._connected_node:
_LOGGER.error(f"Interface {self} cannot be enabled as it is not connected to a Node")
return
if self._connected_node.operating_state != NodeOperatingState.ON:
self._connected_node.sys_log.info(
f"Interface {self} cannot be enabled as the connected Node is not powered on"
)
return
self.enabled = True
self._connected_node.sys_log.info(f"Network Interface {self} enabled")
self.pcap = PacketCapture(hostname=self._connected_node.hostname, interface_num=self.port_num)
AIR_SPACE.add_wireless_interface(self)
def disable(self):
"""Disable the network interface."""
if not self.enabled:
return
self.enabled = False
if self._connected_node:
self._connected_node.sys_log.info(f"Network Interface {self} disabled")
else:
_LOGGER.debug(f"Interface {self} disabled")
AIR_SPACE.remove_wireless_interface(self)
def send_frame(self, frame: Frame) -> bool:
"""
Attempts to send a network frame over the airspace.
This method sends a frame if the network interface is enabled and connected to a wireless airspace. It captures
the frame using PCAP (if available) and transmits it through the airspace. Returns True if the frame is
successfully sent, False otherwise (e.g., if the network interface is disabled).
:param frame: The network frame to be sent.
:return: True if the frame is sent successfully, False if the network interface is disabled.
"""
if self.enabled:
frame.set_sent_timestamp()
self.pcap.capture_outbound(frame)
AIR_SPACE.transmit(frame, self)
return True
# Cannot send Frame as the network interface is not enabled
return False
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the network interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC):
"""
Represents an IP wireless network interface.
This interface operates at both the data link layer (Layer 2) and the network layer (Layer 3) of the OSI model,
specifically tailored for IP-based communication over wireless connections. This abstract class provides a
template for creating specific wireless network interfaces that support Internet Protocol (IP) functionalities.
As this class is a combination of its parent classes without additional attributes or methods, please refer to
the documentation of `WirelessNetworkInterface` and `Layer3Interface` for more details on the supported operations
and functionalities.
The class inherits from:
- `WirelessNetworkInterface`: Providing the functionalities and characteristics of a wireless connection, such as
managing wireless signal transmission, reception, and associated wireless protocols.
- `Layer3Interface`: Enabling network layer capabilities, including IP address assignment, routing, and
potentially, Layer 3 protocols like IPsec.
As an abstract class, `IPWirelessNetworkInterface` does not implement specific methods but ensures that any derived
class provides implementations for the functionalities of both `WirelessNetworkInterface` and `Layer3Interface`.
This setup is ideal for representing network interfaces in devices that require wireless connections and are capable
of IP routing and addressing, such as wireless routers, access points, and wireless end-host devices like
smartphones and laptops.
This class should be extended by concrete classes that define specific behaviors and properties of an IP-capable
wireless network interface.
"""
def model_post_init(self, __context: Any) -> None:
"""
Performs post-initialisation checks to ensure the model's IP configuration is valid.
This method is invoked after the initialisation of a network model object to validate its network settings,
particularly to ensure that the assigned IP address is not a network address. This validation is crucial for
maintaining the integrity of network simulations and avoiding configuration errors that could lead to
unrealistic or incorrect behavior.
:param __context: Contextual information or parameters passed to the method, used for further initializing or
validating the model post-creation.
:raises ValueError: If the IP address is the same as the network address, indicating an incorrect configuration.
"""
if self.ip_network.network_address == self.ip_address:
raise ValueError(f"{self.ip_address}/{self.subnet_mask} must not be a network address")
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
# Get the state from the WiredNetworkInterface
state = WiredNetworkInterface.describe_state(self)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
state["frequency"] = self.frequency
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def enable(self):
"""
Enables this wired network interface and attempts to send a "hello" message to the default gateway.
This method activates the network interface, making it operational for network communications. After enabling,
it tries to initiate a default gateway "hello" process, typically to establish initial connectivity and resolve
the default gateway's MAC address. This step is crucial for ensuring the interface can successfully send data
to and receive data from the network.
The method safely handles cases where the connected node might not have a default gateway set or the
`default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption.
"""
super().enable()
try:
self._connected_node.default_gateway_hello()
except AttributeError:
pass
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass

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

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Optional, Union
from typing import Any, Dict, List, Optional
import matplotlib.pyplot as plt
import networkx as nx
@@ -7,11 +7,11 @@ from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import Router
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.hardware.base import Link, Node, WiredNetworkInterface
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service
@@ -55,8 +55,8 @@ class Network(SimComponent):
for node in self.nodes.values():
node.power_on()
for nic in node.nics.values():
nic.enable()
for network_interface in node.network_interfaces.values():
network_interface.enable()
# Reset software
for software in node.software_manager.software.values():
if isinstance(software, Service):
@@ -141,8 +141,9 @@ class Network(SimComponent):
table.title = "IP Addresses"
for nodes in nodes_type_map.values():
for node in nodes:
for i, port in node.ethernet_port.items():
table.add_row([node.hostname, i, port.ip_address, port.subnet_mask, node.default_gateway])
for i, port in node.network_interface.items():
if hasattr(port, "ip_address"):
table.add_row([node.hostname, i, port.ip_address, port.subnet_mask, node.default_gateway])
print(table)
if links:
@@ -202,8 +203,8 @@ class Network(SimComponent):
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
port_a = link.endpoint_a.port_num
port_b = link.endpoint_b.port_num
state["links"][uuid] = link.describe_state()
state["links"][uuid]["hostname_a"] = hostname_a
state["links"][uuid]["hostname_b"] = hostname_b
@@ -264,18 +265,16 @@ class Network(SimComponent):
self._node_request_manager.remove_request(name=node.hostname)
_LOGGER.info(f"Removed node {node.hostname} from network {self.uuid}")
def connect(
self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs
) -> Optional[Link]:
def connect(self, endpoint_a: WiredNetworkInterface, endpoint_b: WiredNetworkInterface, **kwargs) -> Optional[Link]:
"""
Connect two endpoints on the network by creating a link between their NICs/SwitchPorts.
.. note:: If the nodes owning the endpoints are not already in the network, they are automatically added.
:param endpoint_a: The first endpoint to connect.
:type endpoint_a: Union[NIC, SwitchPort]
:type endpoint_a: WiredNetworkInterface
:param endpoint_b: The second endpoint to connect.
:type endpoint_b: Union[NIC, SwitchPort]
:type endpoint_b: WiredNetworkInterface
:raises RuntimeError: If any validation or runtime checks fail.
"""
node_a: Node = endpoint_a.parent

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

@@ -2,9 +2,9 @@ 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.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
@@ -109,9 +109,9 @@ def create_office_lan(
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])
network.connect(core_switch.network_interface[core_switch_port], switch.network_interface[24])
else:
network.connect(router.ethernet_ports[1], switch.switch_ports[24])
network.connect(router.network_interface[1], switch.network_interface[24])
# Add PCs to the LAN and connect them to switches
for i in range(1, num_pcs + 1):
@@ -125,9 +125,9 @@ def create_office_lan(
# 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])
network.connect(core_switch.network_interface[core_switch_port], switch.network_interface[24])
else:
network.connect(router.ethernet_ports[1], switch.switch_ports[24])
network.connect(router.network_interface[1], switch.network_interface[24])
# Create and add a PC to the network
pc = Computer(
@@ -142,7 +142,7 @@ def create_office_lan(
# 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()
network.connect(switch.network_interface[switch_port], pc.network_interface[1])
switch.network_interface[switch_port].enable()
return network

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

File diff suppressed because it is too large Load Diff

0
src/primaite/simulator/network/hardware/network_interface/__init__.py Normal file
View File

0
src/primaite/simulator/network/hardware/network_interface/wireless/__init__.py Normal file
View File

86
src/primaite/simulator/network/hardware/network_interface/wireless/wireless_access_point.py Normal file
View File

@@ -0,0 +1,86 @@
from typing import Dict
from primaite.simulator.network.hardware.base import (
IPWirelessNetworkInterface,
Layer3Interface,
WirelessNetworkInterface,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
class WirelessAccessPoint(IPWirelessNetworkInterface):
"""
Represents a Wireless Access Point (AP) in a network.
This class models a Wireless Access Point, a device that allows wireless devices to connect to a wired network
using Wi-Fi or other wireless standards. The Wireless Access Point bridges the wireless and wired segments of
the network, allowing wireless devices to communicate with other devices on the network.
As an integral component of wireless networking, a Wireless Access Point provides functionalities for network
management, signal broadcasting, security enforcement, and connection handling. It also possesses Layer 3
capabilities such as IP addressing and subnetting, allowing for network segmentation and routing.
Inherits from:
- WirelessNetworkInterface: Provides basic properties and methods specific to wireless interfaces.
- Layer3Interface: Provides Layer 3 properties like ip_address and subnet_mask, enabling the device to manage
network traffic and routing.
This class can be further specialised or extended to support specific features or standards related to wireless
networking, such as different Wi-Fi versions, frequency bands, or advanced security protocols.
"""
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
# Get the state from the WirelessNetworkInterface
state = WirelessNetworkInterface.describe_state(self)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
# Update the state with NIC-specific information
state.update(
{
"wake_on_lan": self.wake_on_lan,
}
)
return state
def enable(self):
"""Enable the interface."""
pass
def disable(self):
"""Disable the interface."""
pass
def send_frame(self, frame: Frame) -> bool:
"""
Attempts to send a network frame through the interface.
:param frame: The network frame to be sent.
:return: A boolean indicating whether the frame was successfully sent.
"""
pass
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
def __str__(self) -> str:
"""
String representation of the NIC.
:return: A string combining the port number, MAC address and IP address of the NIC.
"""
return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}"

83
src/primaite/simulator/network/hardware/network_interface/wireless/wireless_nic.py Normal file
View File

@@ -0,0 +1,83 @@
from typing import Dict
from primaite.simulator.network.hardware.base import (
IPWirelessNetworkInterface,
Layer3Interface,
WirelessNetworkInterface,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
class WirelessNIC(IPWirelessNetworkInterface):
"""
Represents a Wireless Network Interface Card (Wireless NIC) in a network device.
This class encapsulates the functionalities and attributes of a wireless NIC, combining the characteristics of a
wireless network interface with Layer 3 features. It is capable of connecting to wireless networks, managing
wireless-specific properties such as signal strength and security protocols, and also handling IP-related
functionalities like IP addressing and subnetting.
Inherits from:
- WirelessNetworkInterface: Provides basic properties and methods specific to wireless interfaces.
- Layer3Interface: Provides Layer 3 properties like ip_address and subnet_mask, enabling the device to participate
in IP-based networking.
This class can be extended to include more advanced features or to tailor its behavior for specific types of
wireless networks or protocols.
"""
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
# Get the state from the WirelessNetworkInterface
state = WirelessNetworkInterface.describe_state(self)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
# Update the state with NIC-specific information
state.update(
{
"wake_on_lan": self.wake_on_lan,
}
)
return state
def enable(self):
"""Enable the interface."""
pass
def disable(self):
"""Disable the interface."""
pass
def send_frame(self, frame: Frame) -> bool:
"""
Attempts to send a network frame through the interface.
:param frame: The network frame to be sent.
:return: A boolean indicating whether the frame was successfully sent.
"""
pass
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
def __str__(self) -> str:
"""
String representation of the NIC.
:return: A string combining the port number, MAC address and IP address of the NIC.
"""
return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}"

55
src/primaite/simulator/network/hardware/nodes/computer.py
View File

@@ -1,55 +0,0 @@
from primaite.simulator.network.hardware.base import NIC, Node
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(Node):
"""
A basic Computer class.
Example:
>>> pc_a = Computer(
hostname="pc_a",
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1"
)
>>> pc_a.power_on()
Instances of computer come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"]))
self._install_system_software()
def _install_system_software(self):
"""Install System Software - software that is usually provided with the OS."""
# DNS Client
self.software_manager.install(DNSClient)
# FTP
self.software_manager.install(FTPClient)
# Web Browser
self.software_manager.install(WebBrowser)
super()._install_system_software()

0
src/primaite/simulator/network/hardware/nodes/host/__init__.py Normal file
View File

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

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

385
src/primaite/simulator/network/hardware/nodes/host/host_node.py Normal file
View File

@@ -0,0 +1,385 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from primaite import getLogger
from primaite.simulator.network.hardware.base import IPWiredNetworkInterface, Link, Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.arp.arp import ARP, ARPPacket
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.icmp.icmp import ICMP
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.utils.validators import IPV4Address
_LOGGER = getLogger(__name__)
class HostARP(ARP):
"""
The Host ARP Service.
Extends the ARP service for host-specific functionalities within a network, focusing on resolving and caching
MAC addresses and network interfaces (NICs) based on IP addresses, especially concerning the default gateway.
This specialized ARP service for hosts facilitates efficient network communication by managing ARP entries
and handling ARP requests and replies with additional logic for default gateway processing.
"""
def get_default_gateway_mac_address(self) -> Optional[str]:
"""
Retrieves the MAC address of the default gateway as known from the ARP cache.
:return: The MAC address of the default gateway if present in the ARP cache; otherwise, None.
"""
if self.software_manager.node.default_gateway:
return self.get_arp_cache_mac_address(self.software_manager.node.default_gateway)
def get_default_gateway_network_interface(self) -> Optional[NIC]:
"""
Obtains the network interface card (NIC) associated with the default gateway from the ARP cache.
:return: The NIC associated with the default gateway if it exists in the ARP cache; otherwise, None.
"""
if self.software_manager.node.default_gateway and self.software_manager.node.has_enabled_network_interface:
return self.get_arp_cache_network_interface(self.software_manager.node.default_gateway)
def _get_arp_cache_mac_address(
self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False
) -> Optional[str]:
"""
Internal method to retrieve the MAC address associated with an IP address from the ARP cache.
:param ip_address: The IP address whose MAC address is to be retrieved.
:param is_reattempt: Indicates if this call is a reattempt after a failed initial attempt.
:param is_default_gateway_attempt: Indicates if this call is an attempt to get the default gateway's MAC
address.
:return: The MAC address associated with the IP address if found, otherwise None.
"""
arp_entry = self.arp.get(ip_address)
if arp_entry:
return arp_entry.mac_address
if ip_address == self.software_manager.node.default_gateway:
is_reattempt = True
if not is_reattempt:
self.send_arp_request(ip_address)
return self._get_arp_cache_mac_address(
ip_address=ip_address, is_reattempt=True, is_default_gateway_attempt=is_default_gateway_attempt
)
else:
if self.software_manager.node.default_gateway:
if not is_default_gateway_attempt:
self.send_arp_request(self.software_manager.node.default_gateway)
return self._get_arp_cache_mac_address(
ip_address=self.software_manager.node.default_gateway,
is_reattempt=True,
is_default_gateway_attempt=True,
)
return None
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 for which the MAC address is sought.
:return: The MAC address if available in the ARP cache; otherwise, None.
"""
return self._get_arp_cache_mac_address(ip_address)
def _get_arp_cache_network_interface(
self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False
) -> Optional[NIC]:
"""
Internal method to retrieve the NIC associated with an IP address from the ARP cache.
:param ip_address: The IP address whose NIC is to be retrieved.
:param is_reattempt: Indicates if this call is a reattempt after a failed initial attempt.
:param is_default_gateway_attempt: Indicates if this call is an attempt to get the NIC of the default gateway.
:return: The NIC associated with the IP address if found, otherwise None.
"""
arp_entry = self.arp.get(ip_address)
if arp_entry:
return self.software_manager.node.network_interfaces[arp_entry.network_interface_uuid]
else:
if ip_address == self.software_manager.node.default_gateway:
is_reattempt = True
if not is_reattempt:
self.send_arp_request(ip_address)
return self._get_arp_cache_network_interface(
ip_address=ip_address, is_reattempt=True, is_default_gateway_attempt=is_default_gateway_attempt
)
else:
if self.software_manager.node.default_gateway:
if not is_default_gateway_attempt:
self.send_arp_request(self.software_manager.node.default_gateway)
return self._get_arp_cache_network_interface(
ip_address=self.software_manager.node.default_gateway,
is_reattempt=True,
is_default_gateway_attempt=True,
)
return None
def get_arp_cache_network_interface(self, ip_address: IPv4Address) -> Optional[NIC]:
"""
Retrieves the network interface card (NIC) associated with a given IP address from the ARP cache.
:param ip_address: The IP address for which the associated NIC is sought.
:return: The NIC if available in the ARP cache; otherwise, None.
"""
return self._get_arp_cache_network_interface(ip_address)
def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: NIC):
"""
Processes an ARP request.
Adds a new entry to the ARP cache if the target IP address matches the NIC's IP address and sends an ARP
reply back.
:param arp_packet: The ARP packet containing the request.
:param from_network_interface: The NIC that received the ARP request.
"""
super()._process_arp_request(arp_packet, from_network_interface)
# Unmatched ARP Request
if arp_packet.target_ip_address != from_network_interface.ip_address:
self.sys_log.info(
f"Ignoring ARP request for {arp_packet.target_ip_address}. Current IP address is "
f"{from_network_interface.ip_address}"
)
return
arp_packet = arp_packet.generate_reply(from_network_interface.mac_address)
self.send_arp_reply(arp_packet)
class NIC(IPWiredNetworkInterface):
"""
Represents a Network Interface Card (NIC) in a Host Node.
A NIC is a hardware component that provides a computer or other network device with the ability to connect to a
network. It operates at both Layer 2 (Data Link Layer) and Layer 3 (Network Layer) of the OSI model, meaning it
can interpret both MAC addresses and IP addresses. This class combines the functionalities of
WiredNetworkInterface and Layer3Interface, allowing the NIC to manage physical connections and network layer
addressing.
Inherits from:
- WiredNetworkInterface: Provides properties and methods specific to wired connections, including methods to connect
and disconnect from network links and to manage the enabled/disabled state of the interface.
- Layer3Interface: Provides properties for Layer 3 network configuration, such as IP address and subnet mask.
"""
_connected_link: Optional[Link] = None
"The network link to which the network interface is connected."
wake_on_lan: bool = False
"Indicates if the NIC supports Wake-on-LAN functionality."
def model_post_init(self, __context: Any) -> None:
"""
Performs post-initialisation checks to ensure the model's IP configuration is valid.
This method is invoked after the initialisation of a network model object to validate its network settings,
particularly to ensure that the assigned IP address is not a network address. This validation is crucial for
maintaining the integrity of network simulations and avoiding configuration errors that could lead to
unrealistic or incorrect behavior.
:param __context: Contextual information or parameters passed to the method, used for further initializing or
validating the model post-creation.
:raises ValueError: If the IP address is the same as the network address, indicating an incorrect configuration.
"""
if self.ip_network.network_address == self.ip_address:
raise ValueError(f"{self.ip_address}/{self.subnet_mask} must not be a network address")
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
# Get the state from the IPWiredNetworkInterface
state = super().describe_state()
# Update the state with NIC-specific information
state.update(
{
"wake_on_lan": self.wake_on_lan,
}
)
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def receive_frame(self, frame: Frame) -> bool:
"""
Attempt to receive and process a network frame from the connected Link.
This method processes a frame if the NIC is enabled. It checks the frame's destination and TTL, captures the
frame using PCAP, and forwards it to the connected Node if valid. Returns True if the frame is processed,
False otherwise (e.g., if the NIC is disabled, or TTL expired).
:param frame: The network frame being received.
:return: 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(f"Frame discarded at {self} as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture_inbound(frame)
# If this destination or is broadcast
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_network_interface=self)
return True
return False
def __str__(self) -> str:
"""
String representation of the NIC.
:return: A string combining the port number, MAC address and IP address of the NIC.
"""
return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}"
class HostNode(Node):
"""
Represents a host node in the network.
An end-user device within the network, such as a computer or server, equipped with the capability to initiate and
respond to network communications.
A `HostNode` extends the base `Node` class by incorporating host-specific services and applications, thereby
simulating the functionalities typically expected from a networked end-user device.
**Example**::
>>> pc_a = HostNode(
... hostname="pc_a",
... ip_address="192.168.1.10",
... subnet_mask="255.255.255.0",
... default_gateway="192.168.1.1"
... )
>>> pc_a.power_on()
The host node comes pre-equipped with a range of core functionalities, services, and applications necessary
for engaging in various network operations and tasks.
Core Functionality:
-------------------
* Packet Capture: Monitors and logs network traffic.
* Sys Log: Logs system events and errors.
Services:
---------
* ARP (Address Resolution Protocol) Service: Resolves IP addresses to MAC addresses.
* ICMP (Internet Control Message Protocol) Service: Handles ICMP operations, such as ping requests.
* DNS (Domain Name System) Client: Resolves domain names to IP addresses.
* FTP (File Transfer Protocol) Client: Enables file transfers between the host and FTP servers.
* NTP (Network Time Protocol) Client: Synchronizes the system clock with NTP servers.
Applications:
------------
* Web Browser: Provides web browsing capabilities.
"""
network_interfaces: Dict[str, NIC] = {}
"The Network Interfaces on the node."
network_interface: Dict[int, NIC] = {}
"The NICs on the node by port id."
def __init__(self, ip_address: IPV4Address, subnet_mask: IPV4Address, **kwargs):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
def _install_system_software(self):
"""
Installs the system software and network services typically found on an operating system.
This method equips the host with essential network services and applications, preparing it for various
network-related tasks and operations.
"""
# ARP Service
self.software_manager.install(HostARP)
# ICMP Service
self.software_manager.install(ICMP)
# DNS Client
self.software_manager.install(DNSClient)
# FTP Client
self.software_manager.install(FTPClient)
# NTP Client
self.software_manager.install(NTPClient)
# Web Browser
self.software_manager.install(WebBrowser)
super()._install_system_software()
def default_gateway_hello(self):
"""
Sends a hello message to the default gateway to establish connectivity and resolve the gateway's MAC address.
This method is invoked to ensure the host node can communicate with its default gateway, primarily to confirm
network connectivity and populate the ARP cache with the gateway's MAC address.
"""
if self.operating_state == NodeOperatingState.ON and self.default_gateway:
self.software_manager.arp.get_default_gateway_mac_address()
def receive_frame(self, frame: Frame, from_network_interface: NIC):
"""
Receive a Frame from the connected NIC and process it.
Depending on the protocol, the frame is passed to the appropriate handler such as ARP or ICMP, or up to the
SessionManager if no code manager exists.
:param frame: The Frame being received.
:param from_network_interface: The NIC that received the frame.
"""
super().receive_frame(frame, from_network_interface)
# Check if the destination port is open on the Node
dst_port = None
if frame.tcp:
dst_port = frame.tcp.dst_port
elif frame.udp:
dst_port = frame.udp.dst_port
accept_frame = False
if frame.icmp or dst_port in self.software_manager.get_open_ports():
# accept the frame as the port is open or if it's an ICMP frame
accept_frame = True
# TODO: add internal node firewall check here?
if accept_frame:
self.session_manager.receive_frame(frame, from_network_interface)
else:
self.sys_log.info(f"Ignoring frame from {frame.ip.src_ip_address}")
# TODO: do we need to do anything more here?
pass

12
src/primaite/simulator/network/hardware/nodes/server.py → src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -1,7 +1,7 @@
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
class Server(Computer):
class Server(HostNode):
"""
A basic Server class.
@@ -17,18 +17,14 @@ class Server(Computer):
Instances of Server come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* ARP Service
* ICMP Service
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""

0
src/primaite/simulator/network/hardware/nodes/network/__init__.py Normal file
View File

493
src/primaite/simulator/network/hardware/nodes/network/firewall.py Normal file
View File

@@ -0,0 +1,493 @@
from typing import Dict, Final, Optional, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from primaite.simulator.network.hardware.nodes.network.router import (
AccessControlList,
ACLAction,
Router,
RouterInterface,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.core.sys_log import SysLog
from primaite.utils.validators import IPV4Address
EXTERNAL_PORT_ID: Final[int] = 1
"""The Firewall port ID of the external port."""
INTERNAL_PORT_ID: Final[int] = 2
"""The Firewall port ID of the internal port."""
DMZ_PORT_ID: Final[int] = 3
"""The Firewall port ID of the DMZ port."""
class Firewall(Router):
"""
A Firewall class that extends the functionality of a Router.
The Firewall class acts as a network security system that monitors and controls incoming and outgoing
network traffic based on predetermined security rules. It is an intermediary between internal and external
networks (including DMZ - De-Militarized Zone), ensuring that all inbound and outbound traffic complies with
the security policies.
The Firewall employs Access Control Lists (ACLs) to filter traffic. Both the internal and DMZ ports have both
inbound and outbound ACLs that determine what traffic is allowed to pass.
In addition to the security functions, the Firewall can also perform some routing functions similar to a Router,
forwarding packets between its interfaces based on the destination IP address.
Usage:
To utilise the Firewall class, instantiate it with a hostname and optionally specify sys_log for logging.
Configure the internal, external, and DMZ ports with IP addresses and subnet masks. Define ACL rules to
permit or deny traffic based on your security policies. The Firewall will process frames based on these
rules, determining whether to allow or block traffic at each network interface.
Example:
>>> from primaite.simulator.network.transmission.network_layer import IPProtocol
>>> from primaite.simulator.network.transmission.transport_layer import Port
>>> firewall = Firewall(hostname="Firewall1")
>>> firewall.configure_internal_port(ip_address="192.168.1.1", subnet_mask="255.255.255.0")
>>> firewall.configure_external_port(ip_address="10.0.0.1", subnet_mask="255.255.255.0")
>>> firewall.configure_dmz_port(ip_address="172.16.0.1", subnet_mask="255.255.255.0")
>>> # Permit HTTP traffic to the DMZ
>>> firewall.dmz_inbound_acl.add_rule(
... action=ACLAction.PERMIT,
... protocol=IPProtocol.TCP,
... dst_port=Port.HTTP,
... src_ip_address="0.0.0.0",
... src_wildcard_mask="0.0.0.0",
... dst_ip_address="172.16.0.0",
... dst_wildcard_mask="0.0.0.255"
... )
:ivar str hostname: The Firewall hostname.
"""
internal_inbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing entering the internal network."""
internal_outbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing traffic leaving the internal network."""
dmz_inbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing traffic entering the DMZ."""
dmz_outbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing traffic leaving the DMZ."""
external_inbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing traffic entering from an external network."""
external_outbound_acl: Optional[AccessControlList] = None
"""Access Control List for managing traffic leaving towards an external network."""
def __init__(self, hostname: str, **kwargs):
if not kwargs.get("sys_log"):
kwargs["sys_log"] = SysLog(hostname)
super().__init__(hostname=hostname, num_ports=3, **kwargs)
# Initialise ACLs for internal and dmz interfaces with a default DENY policy
self.internal_inbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY, name=f"{hostname} - Internal Inbound"
)
self.internal_outbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY, name=f"{hostname} - Internal Outbound"
)
self.dmz_inbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY, name=f"{hostname} - DMZ Inbound"
)
self.dmz_outbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY, name=f"{hostname} - DMZ Outbound"
)
# external ACLs should have a default PERMIT policy
self.external_inbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.PERMIT, name=f"{hostname} - External Inbound"
)
self.external_outbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.PERMIT, name=f"{hostname} - External Outbound"
)
self.set_original_state()
def set_original_state(self):
"""Set the original state for the Firewall."""
super().set_original_state()
vals_to_include = {
"internal_port",
"external_port",
"dmz_port",
"internal_inbound_acl",
"internal_outbound_acl",
"dmz_inbound_acl",
"dmz_outbound_acl",
"external_inbound_acl",
"external_outbound_acl",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def describe_state(self) -> Dict:
"""
Describes the current state of the Firewall.
:return: A dictionary representing the current state.
"""
state = super().describe_state()
state.update(
{
"internal_port": self.internal_port.describe_state(),
"external_port": self.external_port.describe_state(),
"dmz_port": self.dmz_port.describe_state(),
"internal_inbound_acl": self.internal_inbound_acl.describe_state(),
"internal_outbound_acl": self.internal_outbound_acl.describe_state(),
"dmz_inbound_acl": self.dmz_inbound_acl.describe_state(),
"dmz_outbound_acl": self.dmz_outbound_acl.describe_state(),
"external_inbound_acl": self.external_inbound_acl.describe_state(),
"external_outbound_acl": self.external_outbound_acl.describe_state(),
}
)
return state
def show(self, markdown: bool = False):
"""
Displays the current configuration of the firewall's network interfaces in a table format.
The table includes information about each port (External, Internal, DMZ), their MAC addresses, IP
configurations, link speeds, and operational status. The output can be formatted as Markdown if specified.
:param markdown: If True, formats the output table in Markdown style. Useful for documentation or reporting
purposes within Markdown-compatible platforms.
"""
table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.hostname} Network Interfaces"
ports = {"External": self.external_port, "Internal": self.internal_port, "DMZ": self.dmz_port}
for port, network_interface in ports.items():
table.add_row(
[
port,
network_interface.mac_address,
f"{network_interface.ip_address}/{network_interface.ip_network.prefixlen}",
network_interface.speed,
"Enabled" if network_interface.enabled else "Disabled",
]
)
print(table)
def show_rules(self, external: bool = True, internal: bool = True, dmz: bool = True, markdown: bool = False):
"""
Prints the configured ACL rules for each specified network zone of the firewall.
This method allows selective viewing of ACL rules applied to external, internal, and DMZ interfaces, providing
a clear overview of the firewall's current traffic filtering policies. Each section can be independently
toggled.
:param external: If True, shows ACL rules for external interfaces.
:param internal: If True, shows ACL rules for internal interfaces.
:param dmz: If True, shows ACL rules for DMZ interfaces.
:param markdown: If True, formats the output in Markdown, enhancing readability in Markdown-compatible viewers.
"""
print(f"{self.hostname} Firewall Rules")
print()
if external:
self.external_inbound_acl.show(markdown)
print()
self.external_outbound_acl.show(markdown)
print()
if internal:
self.internal_inbound_acl.show(markdown)
print()
self.internal_outbound_acl.show(markdown)
print()
if dmz:
self.dmz_inbound_acl.show(markdown)
print()
self.dmz_outbound_acl.show(markdown)
print()
def receive_frame(self, frame: Frame, from_network_interface: RouterInterface):
"""
Receive a frame and process it.
Acts as the primary entry point for all network frames arriving at the Firewall, determining the flow of
traffic based on the source network interface controller (NIC) and applying the appropriate Access Control
List (ACL) rules.
This method categorizes the incoming traffic into three main pathways based on the source NIC: external inbound,
internal outbound, and DMZ (De-Militarized Zone) outbound. It plays a crucial role in enforcing the firewall's
security policies by directing each frame to the corresponding processing method that evaluates it against
specific ACL rules.
Based on the originating NIC:
- Frames from the external port are processed as external inbound traffic, potentially destined for either the
DMZ or the internal network.
- Frames from the internal port are treated as internal outbound traffic, aimed at reaching the external
network or a service within the DMZ.
- Frames from the DMZ port are handled as DMZ outbound traffic, with potential destinations including the
internal network or the external network.
:param frame: The network frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming. Used to
determine the direction of the traffic (inbound or outbound) and the zone (external, internal,
DMZ) it belongs to.
"""
# If the frame comes from the external port, it's considered as external inbound traffic
if from_network_interface == self.external_port:
self._process_external_inbound_frame(frame, from_network_interface)
return
# If the frame comes from the internal port, it's considered as internal outbound traffic
elif from_network_interface == self.internal_port:
self._process_internal_outbound_frame(frame, from_network_interface)
return
# If the frame comes from the DMZ port, it's considered as DMZ outbound traffic
elif from_network_interface == self.dmz_port:
self._process_dmz_outbound_frame(frame, from_network_interface)
return
def _process_external_inbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames arriving from the external network.
Determines the path for frames based on their destination IP addresses and ACL rules for the external inbound
interface. Frames destined for the DMZ or internal network are forwarded accordingly, if allowed by the ACL.
If a frame is permitted by the ACL, it is either passed to the session manager (if applicable) or forwarded to
the appropriate network zone (DMZ/internal). Denied frames are logged and dropped.
:param frame: The frame to be processed, containing network layer and transport layer information.
:param from_network_interface: The interface on the firewall through which the frame was received.
"""
# check if External Inbound ACL Rules permit frame
permitted, rule = self.external_inbound_acl.is_permitted(frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.software_manager.arp.add_arp_cache_entry(
ip_address=frame.ip.src_ip_address,
mac_address=frame.ethernet.src_mac_addr,
network_interface=from_network_interface,
)
if self.check_send_frame_to_session_manager(frame):
# Port is open on this Router so pass Frame up to session manager first
self.session_manager.receive_frame(frame, from_network_interface)
else:
# If the destination IP is within the DMZ network, process the frame as DMZ inbound
if frame.ip.dst_ip_address in self.dmz_port.ip_network:
self._process_dmz_inbound_frame(frame, from_network_interface)
else:
# Otherwise, process the frame as internal inbound
self._process_internal_inbound_frame(frame, from_network_interface)
def _process_external_outbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames that are outbound towards the external network.
:param frame: The frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming.
:param re_attempt: Indicates if the processing is a re-attempt, defaults to False.
"""
# check if External Outbound ACL Rules permit frame
permitted, rule = self.external_outbound_acl.is_permitted(frame=frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.process_frame(frame=frame, from_network_interface=from_network_interface)
def _process_internal_inbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames that are inbound towards the internal LAN.
This method is responsible for handling frames coming from either the external network or the DMZ towards
the internal LAN. It checks the frames against the internal inbound ACL to decide whether to allow or deny
the traffic, and take appropriate actions.
:param frame: The frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming.
:param re_attempt: Indicates if the processing is a re-attempt, defaults to False.
"""
# check if Internal Inbound ACL Rules permit frame
permitted, rule = self.internal_inbound_acl.is_permitted(frame=frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.process_frame(frame=frame, from_network_interface=from_network_interface)
def _process_internal_outbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames that are outbound from the internal network.
This method handles frames that are leaving the internal network. Depending on the destination IP address,
the frame may be forwarded to the DMZ or to the external network.
:param frame: The frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming.
:param re_attempt: Indicates if the processing is a re-attempt, defaults to False.
"""
permitted, rule = self.internal_outbound_acl.is_permitted(frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.software_manager.arp.add_arp_cache_entry(
ip_address=frame.ip.src_ip_address,
mac_address=frame.ethernet.src_mac_addr,
network_interface=from_network_interface,
)
if self.check_send_frame_to_session_manager(frame):
# Port is open on this Router so pass Frame up to session manager first
self.session_manager.receive_frame(frame, from_network_interface)
else:
# If the destination IP is within the DMZ network, process the frame as DMZ inbound
if frame.ip.dst_ip_address in self.dmz_port.ip_network:
self._process_dmz_inbound_frame(frame, from_network_interface)
else:
# If the destination IP is not within the DMZ network, process the frame as external outbound
self._process_external_outbound_frame(frame, from_network_interface)
def _process_dmz_inbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames that are inbound from the DMZ.
This method is responsible for handling frames coming from either the external network or the internal LAN
towards the DMZ. It checks the frames against the DMZ inbound ACL to decide whether to allow or deny the
traffic, and take appropriate actions.
:param frame: The frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming.
:param re_attempt: Indicates if the processing is a re-attempt, defaults to False.
"""
# check if DMZ Inbound ACL Rules permit frame
permitted, rule = self.dmz_inbound_acl.is_permitted(frame=frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.process_frame(frame=frame, from_network_interface=from_network_interface)
def _process_dmz_outbound_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None:
"""
Process frames that are outbound from the DMZ.
This method handles frames originating from the DMZ and determines their appropriate path based on the
destination IP address. It involves checking the DMZ outbound ACL, consulting the ARP cache and the routing
table to find the correct outbound NIC, and then forwarding the frame to either the internal network or the
external network.
:param frame: The frame to be processed.
:param from_network_interface: The network interface controller from which the frame is coming.
:param re_attempt: Indicates if the processing is a re-attempt, defaults to False.
"""
permitted, rule = self.dmz_outbound_acl.is_permitted(frame)
if not permitted:
self.sys_log.info(f"Frame blocked at interface {from_network_interface} by rule {rule}")
return
self.software_manager.arp.add_arp_cache_entry(
ip_address=frame.ip.src_ip_address,
mac_address=frame.ethernet.src_mac_addr,
network_interface=from_network_interface,
)
if self.check_send_frame_to_session_manager(frame):
# Port is open on this Router so pass Frame up to session manager first
self.session_manager.receive_frame(frame, from_network_interface)
else:
# Attempt to get the outbound NIC from the ARP cache using the destination IP address
outbound_nic = self.software_manager.arp.get_arp_cache_network_interface(frame.ip.dst_ip_address)
# If outbound NIC is not found in the ARP cache, consult the routing table to find the best route
if not outbound_nic:
route = self.route_table.find_best_route(frame.ip.dst_ip_address)
if route:
# If a route is found, get the corresponding outbound NIC from the ARP cache using the next-hop IP
# address
outbound_nic = self.software_manager.arp.get_arp_cache_network_interface(route.next_hop_ip_address)
# If an outbound NIC is determined
if outbound_nic:
if outbound_nic == self.external_port:
# If the outbound NIC is the external port, check the frame against the DMZ outbound ACL and
# process it as an external outbound frame
self._process_external_outbound_frame(frame, from_network_interface)
return
elif outbound_nic == self.internal_port:
# If the outbound NIC is the internal port, check the frame against the DMZ outbound ACL and
# process it as an internal inbound frame
self._process_internal_inbound_frame(frame, from_network_interface)
return
# TODO: What to do here? Destination unreachable? Send ICMP back?
return
@property
def external_port(self) -> RouterInterface:
"""
The external port of the firewall.
:return: The external port connecting the firewall to the external network.
"""
return self.network_interface[EXTERNAL_PORT_ID]
@validate_call()
def configure_external_port(self, ip_address: Union[IPV4Address, str], subnet_mask: Union[IPV4Address, str]):
"""
Configure the external port with an IP address and a subnet mask.
Enables the port once configured.
:param ip_address: The IP address to assign to the external port.
:param subnet_mask: The subnet mask to assign to the external port.
"""
# Configure the external port with the specified IP address and subnet mask
self.configure_port(EXTERNAL_PORT_ID, ip_address, subnet_mask)
self.external_port.enable()
@property
def internal_port(self) -> RouterInterface:
"""
The internal port of the firewall.
:return: The external port connecting the firewall to the internal LAN.
"""
return self.network_interface[INTERNAL_PORT_ID]
@validate_call()
def configure_internal_port(self, ip_address: Union[IPV4Address, str], subnet_mask: Union[IPV4Address, str]):
"""
Configure the internal port with an IP address and a subnet mask.
Enables the port once configured.
:param ip_address: The IP address to assign to the internal port.
:param subnet_mask: The subnet mask to assign to the internal port.
"""
self.configure_port(INTERNAL_PORT_ID, ip_address, subnet_mask)
self.internal_port.enable()
@property
def dmz_port(self) -> RouterInterface:
"""
The DMZ port of the firewall.
:return: The external port connecting the firewall to the DMZ.
"""
return self.network_interface[DMZ_PORT_ID]
@validate_call()
def configure_dmz_port(self, ip_address: Union[IPV4Address, str], subnet_mask: Union[IPV4Address, str]):
"""
Configure the DMZ port with an IP address and a subnet mask.
Enables the port once configured.
:param ip_address: The IP address to assign to the DMZ port.
:param subnet_mask: The subnet mask to assign to the DMZ port.
"""
self.configure_port(DMZ_PORT_ID, ip_address, subnet_mask)
self.dmz_port.enable()

30
src/primaite/simulator/network/hardware/nodes/network/network_node.py Normal file
View File

@@ -0,0 +1,30 @@
from abc import abstractmethod
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
class NetworkNode(Node):
"""
Represents an abstract base class for a network node that can receive and process network frames.
This class provides a common interface for network nodes such as routers and switches, defining the essential
behavior that allows these devices to handle incoming network traffic. Implementations of this class must
provide functionality for receiving and processing frames received on their network interfaces.
"""
@abstractmethod
def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface):
"""
Abstract method that must be implemented by subclasses to define how to receive and process frames.
This method is called when a frame is received by a network interface belonging to this node. Subclasses
should implement the logic to process the frame, including examining its contents, making forwarding decisions,
or performing any necessary actions based on the frame's protocol and destination.
:param frame: The network frame that has been received.
:type frame: Frame
:param from_network_interface: The network interface on which the frame was received.
:type from_network_interface: NetworkInterface
"""
pass

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

File diff suppressed because it is too large Load Diff

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

@@ -0,0 +1,200 @@
from __future__ import annotations
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.exceptions import NetworkError
from primaite.simulator.network.hardware.base import Link, WiredNetworkInterface
from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode
from primaite.simulator.network.transmission.data_link_layer import Frame
_LOGGER = getLogger(__name__)
class SwitchPort(WiredNetworkInterface):
"""
Represents a Switch Port.
Switch ports connect devices within the same network. They operate at the data link layer (Layer 2) of the OSI model
and are responsible for receiving and forwarding frames based on MAC addresses. Despite operating at Layer 2,
they are an essential part of network infrastructure, enabling LAN segmentation, bandwidth management, and
the creation of VLANs.
Inherits from:
- WiredNetworkInterface: Provides properties and methods specific to wired connections.
Switch ports typically do not have IP addresses assigned to them as they function at Layer 2, but managed switches
can have management IP addresses for remote management and configuration purposes.
"""
_connected_node: Optional[Switch] = None
"The Switch to which the SwitchPort is connected."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"port_num", "mac_address", "speed", "mtu", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
super().set_original_state()
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
state = super().describe_state()
state.update(
{
"mac_address": self.mac_address,
"speed": self.speed,
"mtu": self.mtu,
"enabled": self.enabled,
}
)
return state
def send_frame(self, frame: Frame) -> bool:
"""
Attempts to send a network frame through the interface.
:param frame: The network frame to be sent.
:return: A boolean indicating whether the frame was successfully sent.
"""
if self.enabled:
self.pcap.capture_outbound(frame)
self._connected_link.transmit_frame(sender_nic=self, frame=frame)
return True
# Cannot send Frame as the SwitchPort is not enabled
return False
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
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_inbound(frame)
self._connected_node.receive_frame(frame=frame, from_network_interface=self)
return True
return False
class Switch(NetworkNode):
"""
A class representing a Layer 2 network switch.
:ivar num_ports: The number of ports on the switch. Default is 24.
"""
num_ports: int = 24
"The number of ports on the switch."
network_interfaces: Dict[str, SwitchPort] = {}
"The SwitchPorts on the Switch."
network_interface: Dict[int, SwitchPort] = {}
"The SwitchPorts on the Switch by port id."
mac_address_table: Dict[str, SwitchPort] = {}
"A MAC address table mapping destination MAC addresses to corresponding SwitchPorts."
def __init__(self, **kwargs):
super().__init__(**kwargs)
if not self.network_interface:
self.network_interface = {i: SwitchPort() for i in range(1, self.num_ports + 1)}
for port_num, port in self.network_interface.items():
port._connected_node = self
port.port_num = port_num
port.parent = self
port.port_num = port_num
def show(self, markdown: bool = False):
"""
Prints a table of the SwitchPorts on the Switch.
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Port", "MAC Address", "Speed", "Status"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.hostname} Switch Ports"
for port_num, port in self.network_interface.items():
table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"])
print(table)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
"""
state = super().describe_state()
state["ports"] = {port_num: port.describe_state() for port_num, port in self.network_interface.items()}
state["num_ports"] = self.num_ports # redundant?
state["mac_address_table"] = {mac: port.port_num for mac, port in self.mac_address_table.items()}
return state
def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort):
"""
Private method to add an entry to the MAC address table.
:param mac_address: MAC address to be added.
:param switch_port: Corresponding SwitchPort object.
"""
mac_table_port = self.mac_address_table.get(mac_address)
if not mac_table_port:
self.mac_address_table[mac_address] = switch_port
self.sys_log.info(f"Added MAC table entry: Port {switch_port.port_num} -> {mac_address}")
else:
if mac_table_port != switch_port:
self.mac_address_table.pop(mac_address)
self.sys_log.info(f"Removed MAC table entry: Port {mac_table_port.port_num} -> {mac_address}")
self._add_mac_table_entry(mac_address, switch_port)
def receive_frame(self, frame: Frame, from_network_interface: SwitchPort):
"""
Forward a frame to the appropriate port based on the destination MAC address.
:param frame: The Frame being received.
:param from_network_interface: The SwitchPort that received the frame.
"""
src_mac = frame.ethernet.src_mac_addr
dst_mac = frame.ethernet.dst_mac_addr
self._add_mac_table_entry(src_mac, from_network_interface)
outgoing_port = self.mac_address_table.get(dst_mac)
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.network_interface.values():
if port.enabled and port != from_network_interface:
port.send_frame(frame)
def disconnect_link_from_port(self, link: Link, port_number: int):
"""
Disconnect a given link from the specified port number on the switch.
:param link: The Link object to be disconnected.
:param port_number: The port number on the switch from where the link should be disconnected.
:raise NetworkError: When an invalid port number is provided or the link does not match the connection.
"""
port = self.network_interface.get(port_number)
if port is None:
msg = f"Invalid port number {port_number} on the switch"
_LOGGER.error(msg)
raise NetworkError(msg)
if port._connected_link != link:
msg = f"The link does not match the connection at port number {port_number}"
_LOGGER.error(msg)
raise NetworkError(msg)
port.disconnect_link()

211
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py Normal file
View File

@@ -0,0 +1,211 @@
from typing import Any, Dict, Union
from pydantic import validate_call
from primaite.simulator.network.airspace import AirSpaceFrequency, IPWirelessNetworkInterface
from primaite.simulator.network.hardware.nodes.network.router import Router, RouterInterface
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.utils.validators import IPV4Address
class WirelessAccessPoint(IPWirelessNetworkInterface):
"""
Represents a Wireless Access Point (AP) in a network.
This class models a Wireless Access Point, a device that allows wireless devices to connect to a wired network
using Wi-Fi or other wireless standards. The Wireless Access Point bridges the wireless and wired segments of
the network, allowing wireless devices to communicate with other devices on the network.
As an integral component of wireless networking, a Wireless Access Point provides functionalities for network
management, signal broadcasting, security enforcement, and connection handling. It also possesses Layer 3
capabilities such as IP addressing and subnetting, allowing for network segmentation and routing.
Inherits from:
- WirelessNetworkInterface: Provides basic properties and methods specific to wireless interfaces.
- Layer3Interface: Provides Layer 3 properties like ip_address and subnet_mask, enabling the device to manage
network traffic and routing.
This class can be further specialised or extended to support specific features or standards related to wireless
networking, such as different Wi-Fi versions, frequency bands, or advanced security protocols.
"""
def model_post_init(self, __context: Any) -> None:
"""
Performs post-initialisation checks to ensure the model's IP configuration is valid.
This method is invoked after the initialisation of a network model object to validate its network settings,
particularly to ensure that the assigned IP address is not a network address. This validation is crucial for
maintaining the integrity of network simulations and avoiding configuration errors that could lead to
unrealistic or incorrect behavior.
:param __context: Contextual information or parameters passed to the method, used for further initializing or
validating the model post-creation.
:raises ValueError: If the IP address is the same as the network address, indicating an incorrect configuration.
"""
if self.ip_network.network_address == self.ip_address:
raise ValueError(f"{self.ip_address}/{self.subnet_mask} must not be a network address")
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
:rtype: Dict
"""
return super().describe_state()
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the interface.
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
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_inbound(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_network_interface=self)
return True
return False
def __str__(self) -> str:
"""
String representation of the NIC.
:return: A string combining the port number, MAC address and IP address of the NIC.
"""
return f"Port {self.port_num}: {self.mac_address}/{self.ip_address} ({self.frequency})"
class WirelessRouter(Router):
"""
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.
This class represents a network device that performs routing functions similar to a traditional router but also
includes the functionality of a wireless access point. This allows the WirelessRouter to not only direct traffic
between wired networks but also to manage and facilitate wireless network connections.
A WirelessRouter is instantiated and configured with both wired and wireless interfaces. The wired interfaces are
managed similarly to those in a standard Router, while the wireless interfaces require additional configuration
specific to wireless settings, such as setting the frequency band (e.g., 2.4 GHz or 5 GHz for Wi-Fi).
The WirelessRouter facilitates creating a network environment where devices can be interconnected via both
Ethernet (wired) and Wi-Fi (wireless), making it an essential component for simulating more complex and realistic
network topologies within PrimAITE.
Example:
>>> wireless_router = WirelessRouter(hostname="wireless_router_1")
>>> wireless_router.configure_router_interface(
... ip_address="192.168.1.1",
... subnet_mask="255.255.255.0"
... )
>>> wireless_router.configure_wireless_access_point(
... ip_address="10.10.10.1",
... subnet_mask="255.255.255.0"
... frequency=AirSpaceFrequency.WIFI_2_4
... )
"""
network_interfaces: Dict[str, Union[RouterInterface, WirelessAccessPoint]] = {}
network_interface: Dict[int, Union[RouterInterface, WirelessAccessPoint]] = {}
def __init__(self, hostname: str, **kwargs):
super().__init__(hostname=hostname, num_ports=0, **kwargs)
self.connect_nic(WirelessAccessPoint(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
self.connect_nic(RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
self.set_original_state()
@property
def wireless_access_point(self) -> WirelessAccessPoint:
"""
Retrieves the wireless access point interface associated with this wireless router.
This property provides direct access to the WirelessAccessPoint interface of the router, facilitating wireless
communications. Specifically, it returns the interface configured on port 1, dedicated to establishing and
managing wireless network connections. This interface is essential for enabling wireless connectivity,
allowing devices within connect to the network wirelessly.
:return: The WirelessAccessPoint instance representing the wireless connection interface on port 1 of the
wireless router.
"""
return self.network_interface[1]
@validate_call()
def configure_wireless_access_point(
self,
ip_address: IPV4Address,
subnet_mask: IPV4Address,
frequency: AirSpaceFrequency = AirSpaceFrequency.WIFI_2_4,
):
"""
Configures a wireless access point (WAP).
Sets its IP address, subnet mask, and operating frequency. This method ensures the wireless access point is
properly set up to manage wireless communication over the specified frequency band.
The method first disables the WAP to safely apply configuration changes. After configuring the IP settings,
it sets the WAP to operate on the specified frequency band and then re-enables the WAP for operation.
:param ip_address: The IP address to be assigned to the wireless access point.
:param subnet_mask: The subnet mask associated with the IP address
:param frequency: The operating frequency of the wireless access point, defined by the AirSpaceFrequency
enum. This determines the frequency band (e.g., 2.4 GHz or 5 GHz) the access point will use for wireless
communication. Default is AirSpaceFrequency.WIFI_2_4.
"""
self.wireless_access_point.disable() # Temporarily disable the WAP for reconfiguration
network_interface = self.network_interface[1]
network_interface.ip_address = ip_address
network_interface.subnet_mask = subnet_mask
self.sys_log.info(f"Configured WAP {network_interface}")
self.set_original_state()
self.wireless_access_point.frequency = frequency # Set operating frequency
self.wireless_access_point.enable() # Re-enable the WAP with new settings
@property
def router_interface(self) -> RouterInterface:
"""
Retrieves the router interface associated with this wireless router.
This property provides access to the router interface configured for wired connections. It specifically
returns the interface configured on port 2, which is reserved for wired LAN/WAN connections.
:return: The RouterInterface instance representing the wired LAN/WAN connection on port 2 of the wireless
router.
"""
return self.network_interface[2]
@validate_call()
def configure_router_interface(
self,
ip_address: IPV4Address,
subnet_mask: IPV4Address,
):
"""
Configures a router interface.
Sets its IP address and subnet mask.
The method first disables the router interface to safely apply configuration changes. After configuring the IP
settings, it re-enables the router interface for operation.
:param ip_address: The IP address to be assigned to the router interface.
:param subnet_mask: The subnet mask associated with the IP address
"""
self.router_interface.disable() # Temporarily disable the router interface for reconfiguration
super().configure_port(port=2, ip_address=ip_address, subnet_mask=subnet_mask) # Set IP configuration
self.router_interface.enable() # Re-enable the router interface with new settings
def configure_port(self, port: int, ip_address: Union[IPV4Address, str], subnet_mask: Union[IPV4Address, str]):
"""Not Implemented."""
raise NotImplementedError(
"Please use the 'configure_wireless_access_point' and 'configure_router_interface' functions."
)

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

File diff suppressed because it is too large Load Diff

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

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

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

@@ -1,11 +1,11 @@
from ipaddress import IPv4Address
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import NIC, NodeOperatingState
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.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.network.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
@@ -40,13 +40,13 @@ def client_server_routed() -> Network:
# Switch 1
switch_1 = Switch(hostname="switch_1", num_ports=6)
switch_1.power_on()
network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6])
network.connect(endpoint_a=router_1.network_interface[1], endpoint_b=switch_1.network_interface[6])
router_1.enable_port(1)
# Switch 2
switch_2 = Switch(hostname="switch_2", num_ports=6)
switch_2.power_on()
network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6])
network.connect(endpoint_a=router_1.network_interface[2], endpoint_b=switch_2.network_interface[6])
router_1.enable_port(2)
# Client 1
@@ -55,10 +55,10 @@ def client_server_routed() -> Network:
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.2.1",
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
client_1.power_on()
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
# Server 1
server_1 = Server(
@@ -66,10 +66,10 @@ def client_server_routed() -> Network:
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
server_1.power_on()
network.connect(endpoint_b=server_1.ethernet_port[1], endpoint_a=switch_1.switch_ports[1])
network.connect(endpoint_b=server_1.network_interface[1], endpoint_a=switch_1.network_interface[1])
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
@@ -118,21 +118,21 @@ def arcd_uc2_network() -> Network:
network = Network()
# Router 1
router_1 = Router(hostname="router_1", num_ports=5, operating_state=NodeOperatingState.ON)
router_1 = Router(hostname="router_1", num_ports=5, start_up_duration=0)
router_1.power_on()
router_1.configure_port(port=1, ip_address="192.168.1.1", subnet_mask="255.255.255.0")
router_1.configure_port(port=2, ip_address="192.168.10.1", subnet_mask="255.255.255.0")
# Switch 1
switch_1 = Switch(hostname="switch_1", num_ports=8, operating_state=NodeOperatingState.ON)
switch_1 = Switch(hostname="switch_1", num_ports=8, start_up_duration=0)
switch_1.power_on()
network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[8])
network.connect(endpoint_a=router_1.network_interface[1], endpoint_b=switch_1.network_interface[8])
router_1.enable_port(1)
# Switch 2
switch_2 = Switch(hostname="switch_2", num_ports=8, operating_state=NodeOperatingState.ON)
switch_2 = Switch(hostname="switch_2", num_ports=8, start_up_duration=0)
switch_2.power_on()
network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[8])
network.connect(endpoint_a=router_1.network_interface[2], endpoint_b=switch_2.network_interface[8])
router_1.enable_port(2)
# Client 1
@@ -142,10 +142,10 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.10.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
client_1.power_on()
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
db_manipulation_bot.configure(
@@ -162,12 +162,12 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.10.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
client_2.power_on()
web_browser = client_2.software_manager.software.get("WebBrowser")
web_browser.target_url = "http://arcd.com/users/"
network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2])
network.connect(endpoint_b=client_2.network_interface[1], endpoint_a=switch_2.network_interface[2])
# Domain Controller
domain_controller = Server(
@@ -175,12 +175,12 @@ def arcd_uc2_network() -> Network:
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
domain_controller.power_on()
domain_controller.software_manager.install(DNSServer)
network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1])
network.connect(endpoint_b=domain_controller.network_interface[1], endpoint_a=switch_1.network_interface[1])
# Database Server
database_server = Server(
@@ -189,10 +189,10 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
database_server.power_on()
network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3])
network.connect(endpoint_b=database_server.network_interface[1], endpoint_a=switch_1.network_interface[3])
ddl = """
CREATE TABLE IF NOT EXISTS user (
@@ -263,14 +263,14 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])
network.connect(endpoint_b=web_server.network_interface[1], endpoint_a=switch_1.network_interface[2])
database_client.run()
database_client.connect()
@@ -278,7 +278,7 @@ def arcd_uc2_network() -> Network:
# register the web_server to a domain
dns_server_service: DNSServer = domain_controller.software_manager.software.get("DNSServer") # noqa
dns_server_service.dns_register("arcd.com", web_server.ip_address)
dns_server_service.dns_register("arcd.com", web_server.network_interface[1].ip_address)
# Backup Server
backup_server = Server(
@@ -287,11 +287,11 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
backup_server.power_on()
backup_server.software_manager.install(FTPServer)
network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4])
network.connect(endpoint_b=backup_server.network_interface[1], endpoint_a=switch_1.network_interface[4])
# Security Suite
security_suite = Server(
@@ -300,12 +300,12 @@ def arcd_uc2_network() -> Network:
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
dns_server=IPv4Address("192.168.1.10"),
operating_state=NodeOperatingState.ON,
start_up_duration=0,
)
security_suite.power_on()
network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7])
network.connect(endpoint_b=security_suite.network_interface[1], endpoint_a=switch_1.network_interface[7])
security_suite.connect_nic(NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0"))
network.connect(endpoint_b=security_suite.ethernet_port[2], endpoint_a=switch_2.switch_ports[7])
network.connect(endpoint_b=security_suite.network_interface[2], endpoint_a=switch_2.network_interface[7])
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)

5
src/primaite/simulator/network/protocols/arp.py
View File

@@ -13,11 +13,12 @@ class ARPEntry(BaseModel):
Represents an entry in the ARP cache.
:param mac_address: The MAC address associated with the IP address.
:param nic: The NIC through which the NIC with the IP address is reachable.
:param network_interface_uuid: The UIId of the Network Interface through which the NIC with the IP address is
reachable.
"""
mac_address: str
nic_uuid: str
network_interface_uuid: str
class ARPPacket(DataPacket):

114
src/primaite/simulator/network/protocols/icmp.py Normal file
View File

@@ -0,0 +1,114 @@
import secrets
from enum import Enum
from typing import Union
from pydantic import BaseModel, field_validator, validate_call
from pydantic_core.core_schema import FieldValidationInfo
from primaite import getLogger
_LOGGER = getLogger(__name__)
class ICMPType(Enum):
"""Enumeration of common ICMP (Internet Control Message Protocol) types."""
ECHO_REPLY = 0
"Echo Reply message."
DESTINATION_UNREACHABLE = 3
"Destination Unreachable."
REDIRECT = 5
"Redirect."
ECHO_REQUEST = 8
"Echo Request (ping)."
ROUTER_ADVERTISEMENT = 9
"Router Advertisement."
ROUTER_SOLICITATION = 10
"Router discovery/selection/solicitation."
TIME_EXCEEDED = 11
"Time Exceeded."
TIMESTAMP_REQUEST = 13
"Timestamp Request."
TIMESTAMP_REPLY = 14
"Timestamp Reply."
@validate_call
def get_icmp_type_code_description(icmp_type: ICMPType, icmp_code: int) -> Union[str, None]:
"""
Maps ICMPType and code pairings to their respective description.
:param icmp_type: An ICMPType.
:param icmp_code: An icmp code.
:return: The icmp type and code pairing description if it exists, otherwise returns None.
"""
icmp_code_descriptions = {
ICMPType.ECHO_REPLY: {0: "Echo reply"},
ICMPType.DESTINATION_UNREACHABLE: {
0: "Destination network unreachable",
1: "Destination host unreachable",
2: "Destination protocol unreachable",
3: "Destination port unreachable",
4: "Fragmentation required",
5: "Source route failed",
6: "Destination network unknown",
7: "Destination host unknown",
8: "Source host isolated",
9: "Network administratively prohibited",
10: "Host administratively prohibited",
11: "Network unreachable for ToS",
12: "Host unreachable for ToS",
13: "Communication administratively prohibited",
14: "Host Precedence Violation",
15: "Precedence cutoff in effect",
},
ICMPType.REDIRECT: {
0: "Redirect Datagram for the Network",
1: "Redirect Datagram for the Host",
},
ICMPType.ECHO_REQUEST: {0: "Echo request"},
ICMPType.ROUTER_ADVERTISEMENT: {0: "Router Advertisement"},
ICMPType.ROUTER_SOLICITATION: {0: "Router discovery/selection/solicitation"},
ICMPType.TIME_EXCEEDED: {0: "TTL expired in transit", 1: "Fragment reassembly time exceeded"},
ICMPType.TIMESTAMP_REQUEST: {0: "Timestamp Request"},
ICMPType.TIMESTAMP_REPLY: {0: "Timestamp reply"},
}
return icmp_code_descriptions[icmp_type].get(icmp_code)
class ICMPPacket(BaseModel):
"""Models an ICMP Packet."""
icmp_type: ICMPType = ICMPType.ECHO_REQUEST
"ICMP Type."
icmp_code: int = 0
"ICMP Code."
identifier: int
"ICMP identifier (16 bits randomly generated)."
sequence: int = 0
"ICMP message sequence number."
def __init__(self, **kwargs):
if not kwargs.get("identifier"):
kwargs["identifier"] = secrets.randbits(16)
super().__init__(**kwargs)
@field_validator("icmp_code") # noqa
@classmethod
def _icmp_type_must_have_icmp_code(cls, v: int, info: FieldValidationInfo) -> int:
"""Validates the icmp_type and icmp_code."""
icmp_type = info.data["icmp_type"]
if get_icmp_type_code_description(icmp_type, v):
return v
msg = f"No Matching ICMP code for type:{icmp_type.name}, code:{v}"
_LOGGER.error(msg)
raise ValueError(msg)
def code_description(self) -> str:
"""The icmp_code description."""
description = get_icmp_type_code_description(self.icmp_type, self.icmp_code)
if description:
return description
msg = f"No Matching ICMP code for type:{self.icmp_type.name}, code:{self.icmp_code}"
_LOGGER.error(msg)
raise ValueError(msg)

8
src/primaite/simulator/network/transmission/data_link_layer.py
View File

@@ -4,9 +4,9 @@ from typing import Any, Optional
from pydantic import BaseModel
from primaite import getLogger
from primaite.simulator.network.protocols.arp import ARPPacket
from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.protocols.packet import DataPacket
from primaite.simulator.network.transmission.network_layer import ICMPPacket, IPPacket, IPProtocol
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.primaite_layer import PrimaiteHeader
from primaite.simulator.network.transmission.transport_layer import TCPHeader, UDPHeader
from primaite.simulator.network.utils import convert_bytes_to_megabits
@@ -73,7 +73,7 @@ class Frame(BaseModel):
msg = "Cannot build a Frame using the TCP IP Protocol without a TCPHeader"
_LOGGER.error(msg)
raise ValueError(msg)
if kwargs["ip"].protocol == IPProtocol.UDP and not kwargs.get("UDP"):
if kwargs["ip"].protocol == IPProtocol.UDP and not kwargs.get("udp"):
msg = "Cannot build a Frame using the UDP IP Protocol without a UDPHeader"
_LOGGER.error(msg)
raise ValueError(msg)
@@ -95,8 +95,6 @@ class Frame(BaseModel):
"UDP header."
icmp: Optional[ICMPPacket] = None
"ICMP header."
arp: Optional[ARPPacket] = None
"ARP packet."
primaite: PrimaiteHeader
"PrimAITE header."
payload: Optional[Any] = None

122
src/primaite/simulator/network/transmission/network_layer.py
View File

@@ -1,12 +1,9 @@
import secrets
from enum import Enum
from ipaddress import IPv4Address
from typing import Union
from pydantic import BaseModel, field_validator, validate_call
from pydantic_core.core_schema import FieldValidationInfo
from pydantic import BaseModel
from primaite import getLogger
from primaite.utils.validators import IPV4Address
_LOGGER = getLogger(__name__)
@@ -54,110 +51,6 @@ class Precedence(Enum):
"Highest priority level, used for the most critical network control messages, such as routing protocol hellos."
class ICMPType(Enum):
"""Enumeration of common ICMP (Internet Control Message Protocol) types."""
ECHO_REPLY = 0
"Echo Reply message."
DESTINATION_UNREACHABLE = 3
"Destination Unreachable."
REDIRECT = 5
"Redirect."
ECHO_REQUEST = 8
"Echo Request (ping)."
ROUTER_ADVERTISEMENT = 10
"Router Advertisement."
ROUTER_SOLICITATION = 11
"Router discovery/selection/solicitation."
TIME_EXCEEDED = 11
"Time Exceeded."
TIMESTAMP_REQUEST = 13
"Timestamp Request."
TIMESTAMP_REPLY = 14
"Timestamp Reply."
@validate_call
def get_icmp_type_code_description(icmp_type: ICMPType, icmp_code: int) -> Union[str, None]:
"""
Maps ICMPType and code pairings to their respective description.
:param icmp_type: An ICMPType.
:param icmp_code: An icmp code.
:return: The icmp type and code pairing description if it exists, otherwise returns None.
"""
icmp_code_descriptions = {
ICMPType.ECHO_REPLY: {0: "Echo reply"},
ICMPType.DESTINATION_UNREACHABLE: {
0: "Destination network unreachable",
1: "Destination host unreachable",
2: "Destination protocol unreachable",
3: "Destination port unreachable",
4: "Fragmentation required",
5: "Source route failed",
6: "Destination network unknown",
7: "Destination host unknown",
8: "Source host isolated",
9: "Network administratively prohibited",
10: "Host administratively prohibited",
11: "Network unreachable for ToS",
12: "Host unreachable for ToS",
13: "Communication administratively prohibited",
14: "Host Precedence Violation",
15: "Precedence cutoff in effect",
},
ICMPType.REDIRECT: {
0: "Redirect Datagram for the Network",
1: "Redirect Datagram for the Host",
},
ICMPType.ECHO_REQUEST: {0: "Echo request"},
ICMPType.ROUTER_ADVERTISEMENT: {0: "Router Advertisement"},
ICMPType.ROUTER_SOLICITATION: {0: "Router discovery/selection/solicitation"},
ICMPType.TIME_EXCEEDED: {0: "TTL expired in transit", 1: "Fragment reassembly time exceeded"},
ICMPType.TIMESTAMP_REQUEST: {0: "Timestamp Request"},
ICMPType.TIMESTAMP_REPLY: {0: "Timestamp reply"},
}
return icmp_code_descriptions[icmp_type].get(icmp_code)
class ICMPPacket(BaseModel):
"""Models an ICMP Packet."""
icmp_type: ICMPType = ICMPType.ECHO_REQUEST
"ICMP Type."
icmp_code: int = 0
"ICMP Code."
identifier: int
"ICMP identifier (16 bits randomly generated)."
sequence: int = 0
"ICMP message sequence number."
def __init__(self, **kwargs):
if not kwargs.get("identifier"):
kwargs["identifier"] = secrets.randbits(16)
super().__init__(**kwargs)
@field_validator("icmp_code") # noqa
@classmethod
def _icmp_type_must_have_icmp_code(cls, v: int, info: FieldValidationInfo) -> int:
"""Validates the icmp_type and icmp_code."""
icmp_type = info.data["icmp_type"]
if get_icmp_type_code_description(icmp_type, v):
return v
msg = f"No Matching ICMP code for type:{icmp_type.name}, code:{v}"
_LOGGER.error(msg)
raise ValueError(msg)
def code_description(self) -> str:
"""The icmp_code description."""
description = get_icmp_type_code_description(self.icmp_type, self.icmp_code)
if description:
return description
msg = f"No Matching ICMP code for type:{self.icmp_type.name}, code:{self.icmp_code}"
_LOGGER.error(msg)
raise ValueError(msg)
class IPPacket(BaseModel):
"""
Represents the IP layer of a network frame.
@@ -180,9 +73,9 @@ class IPPacket(BaseModel):
... )
"""
src_ip_address: IPv4Address
src_ip_address: IPV4Address
"Source IP address."
dst_ip_address: IPv4Address
dst_ip_address: IPV4Address
"Destination IP address."
protocol: IPProtocol = IPProtocol.TCP
"IPProtocol."
@@ -190,10 +83,3 @@ class IPPacket(BaseModel):
"Time to Live (TTL) for the packet."
precedence: Precedence = Precedence.ROUTINE
"Precedence level for Quality of Service (default is Precedence.ROUTINE)."
def __init__(self, **kwargs):
if not isinstance(kwargs["src_ip_address"], IPv4Address):
kwargs["src_ip_address"] = IPv4Address(kwargs["src_ip_address"])
if not isinstance(kwargs["dst_ip_address"], IPv4Address):
kwargs["dst_ip_address"] = IPv4Address(kwargs["dst_ip_address"])
super().__init__(**kwargs)

2
src/primaite/simulator/network/transmission/transport_layer.py
View File

@@ -7,6 +7,8 @@ from pydantic import BaseModel
class Port(Enum):
"""Enumeration of common known TCP/UDP ports used by protocols for operation of network applications."""
NONE = 0
"Place holder for a non-port."
WOL = 9
"Wake-on-Lan (WOL) - Used to turn or awaken a computer from sleep mode by a network message."
FTP_DATA = 20

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

@@ -23,6 +23,7 @@ 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):
@@ -59,9 +60,10 @@ class DatabaseClient(Application):
if not connection_id:
connection_id = str(uuid4())
return self._connect(
self.connected = self._connect(
server_ip_address=self.server_ip_address, password=self.server_password, connection_id=connection_id
)
return self.connected
def _connect(
self,
@@ -133,6 +135,7 @@ class DatabaseClient(Application):
self.sys_log.info(
f"{self.name}: DatabaseClient disconnected connection {connection_id} from {self.server_ip_address}"
)
self.connected = False
def _query(self, sql: str, query_id: str, connection_id: str, is_reattempt: bool = False) -> bool:
"""

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

@@ -21,7 +21,7 @@ class PacketCapture:
The PCAPs are logged to: <simulation output directory>/<hostname>/<hostname>_<ip address>_pcap.log
"""
def __init__(self, hostname: str, ip_address: Optional[str] = None, switch_port_number: Optional[int] = None):
def __init__(self, hostname: str, ip_address: Optional[str] = None, interface_num: Optional[int] = None):
"""
Initialize the PacketCapture process.
@@ -32,19 +32,20 @@ class PacketCapture:
"The hostname for which PCAP logs are being recorded."
self.ip_address: str = ip_address
"The IP address associated with the PCAP logs."
self.switch_port_number = switch_port_number
"The SwitchPort number."
self.interface_num = interface_num
"The interface num on the Node."
self.inbound_logger = None
self.outbound_logger = None
self.current_episode: int = 1
self.setup_logger()
self.setup_logger(outbound=False)
self.setup_logger(outbound=True)
def setup_logger(self):
def setup_logger(self, outbound: bool = False):
"""Set up the logger configuration."""
if not SIM_OUTPUT.save_pcap_logs:
return
log_path = self._get_log_path()
log_path = self._get_log_path(outbound)
file_handler = logging.FileHandler(filename=log_path)
file_handler.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs
@@ -52,11 +53,17 @@ class PacketCapture:
log_format = "%(message)s"
file_handler.setFormatter(logging.Formatter(log_format))
self.logger = logging.getLogger(self._logger_name)
self.logger.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs
self.logger.addHandler(file_handler)
if outbound:
self.outbound_logger = logging.getLogger(self._get_logger_name(outbound))
logger = self.outbound_logger
else:
self.inbound_logger = logging.getLogger(self._get_logger_name(outbound))
logger = self.inbound_logger
self.logger.addFilter(_JSONFilter())
logger.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs
logger.addHandler(file_handler)
logger.addFilter(_JSONFilter())
def read(self) -> List[Dict[str, Any]]:
"""
@@ -70,27 +77,34 @@ class PacketCapture:
frames.append(json.loads(line.rstrip()))
return frames
@property
def _logger_name(self) -> str:
def _get_logger_name(self, outbound: bool = False) -> str:
"""Get PCAP the logger name."""
if self.ip_address:
return f"{self.hostname}_{self.ip_address}_pcap"
if self.switch_port_number:
return f"{self.hostname}_port-{self.switch_port_number}_pcap"
return f"{self.hostname}_pcap"
return f"{self.hostname}_{self.ip_address}_{'outbound' if outbound else 'inbound'}_pcap"
if self.interface_num:
return f"{self.hostname}_port-{self.interface_num}_{'outbound' if outbound else 'inbound'}_pcap"
return f"{self.hostname}_{'outbound' if outbound else 'inbound'}_pcap"
def _get_log_path(self) -> Path:
def _get_log_path(self, outbound: bool = False) -> Path:
"""Get the path for the log file."""
root = SIM_OUTPUT.path / f"episode_{self.current_episode}" / self.hostname
root.mkdir(exist_ok=True, parents=True)
return root / f"{self._logger_name}.log"
return root / f"{self._get_logger_name(outbound)}.log"
def capture(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;(
def capture_inbound(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;(
"""
Capture a Frame and log it.
Capture an inbound Frame and log it.
:param frame: The PCAP frame to capture.
"""
if SIM_OUTPUT.save_pcap_logs:
msg = frame.model_dump_json()
self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL
msg = frame.model_dump_json()
self.inbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL
def capture_outbound(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;(
"""
Capture an outbound Frame and log it.
:param frame: The PCAP frame to capture.
"""
msg = frame.model_dump_json()
self.outbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL

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

@@ -6,12 +6,14 @@ from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.core import SimComponent
from primaite.simulator.network.protocols.arp import ARPPacket
from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader, UDPHeader
if TYPE_CHECKING:
from primaite.simulator.network.hardware.base import ARPCache
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.simulator.system.core.sys_log import SysLog
@@ -73,14 +75,14 @@ class SessionManager:
:param arp_cache: A reference to the ARP cache component.
"""
def __init__(self, sys_log: SysLog, arp_cache: "ARPCache"):
def __init__(self, sys_log: SysLog):
self.sessions_by_key: Dict[
Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]], Session
] = {}
self.sessions_by_uuid: Dict[str, Session] = {}
self.sys_log: SysLog = sys_log
self.software_manager: SoftwareManager = None # Noqa
self.arp_cache: "ARPCache" = arp_cache
self.node: Node = None # noqa
def describe_state(self) -> Dict:
"""
@@ -138,13 +140,130 @@ class SessionManager:
dst_port = None
return protocol, with_ip_address, src_port, dst_port
def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional["NetworkInterface"]:
"""
Resolves the appropriate outbound network interface for a given destination IP address.
This method determines the most suitable network interface for sending a packet to the specified
destination IP address. It considers only enabled network interfaces and checks if the destination
IP address falls within the subnet of each interface. If no suitable local network interface is found,
the method defaults to using the network interface associated with the default gateway.
The search process prioritises local network interfaces based on the IP network to which they belong.
If the destination IP address does not match any local subnet, the method assumes that the destination
is outside the local network and hence, routes the packet through the default gateway's network interface.
:param dst_ip_address: The destination IP address for which the outbound interface is to be resolved.
:type dst_ip_address: IPv4Address
:return: The network interface through which the packet should be sent to reach the destination IP address,
or the default gateway's network interface if the destination is not within any local subnet.
:rtype: Optional["NetworkInterface"]
"""
for network_interface in self.node.network_interfaces.values():
if dst_ip_address in network_interface.ip_network and network_interface.enabled:
return network_interface
return self.software_manager.arp.get_default_gateway_network_interface()
def resolve_outbound_transmission_details(
self,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
protocol: Optional[IPProtocol] = None,
session_id: Optional[str] = None,
) -> Tuple[
Optional["NetworkInterface"],
Optional[str],
IPv4Address,
Optional[Port],
Optional[Port],
Optional[IPProtocol],
bool,
]:
"""
Resolves the necessary details for outbound transmission based on the provided parameters.
This method determines whether the payload should be broadcast or unicast based on the destination IP address
and resolves the outbound network interface and destination MAC address accordingly.
The method first checks if `session_id` is provided and uses the session details if available. For broadcast
transmissions, it finds a suitable network interface and uses a broadcast MAC address. For unicast
transmissions, it attempts to resolve the destination MAC address using ARP and finds the appropriate
outbound network interface. If the destination IP address is outside the local network and no specific MAC
address is resolved, it uses the default gateway for the transmission.
:param dst_ip_address: The destination IP address or network. If an IPv4Network is provided, the method
treats the transmission as a broadcast to that network. Optional.
:type dst_ip_address: Optional[Union[IPv4Address, IPv4Network]]
:param src_port: The source port number for the transmission. Optional.
:type src_port: Optional[Port]
:param dst_port: The destination port number for the transmission. Optional.
:type dst_port: Optional[Port]
:param protocol: The IP protocol to be used for the transmission. Optional.
:type protocol: Optional[IPProtocol]
:param session_id: The session ID associated with the transmission. If provided, the session details override
other parameters. Optional.
:type session_id: Optional[str]
:return: A tuple containing the resolved outbound network interface, destination MAC address, destination IP
address, source port, destination port, protocol, and a boolean indicating whether the transmission is a
broadcast.
:rtype: Tuple[Optional["NetworkInterface"], Optional[str], IPv4Address, Optional[Port], Optional[Port],
Optional[IPProtocol], bool]
"""
if dst_ip_address and not isinstance(dst_ip_address, (IPv4Address, IPv4Network)):
dst_ip_address = IPv4Address(dst_ip_address)
is_broadcast = False
outbound_network_interface = 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 = session.with_ip_address
protocol = session.protocol
src_port = session.src_port
dst_port = session.dst_port
# Determine if the payload is for broadcast or unicast
# 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 network_interface in self.node.network_interfaces.values():
if dst_ip_address in network_interface.ip_network and network_interface.enabled:
dst_mac_address = "ff:ff:ff:ff:ff:ff"
outbound_network_interface = network_interface
break
else:
# Resolve MAC address for unicast transmission
use_default_gateway = True
for network_interface in self.node.network_interfaces.values():
if dst_ip_address in network_interface.ip_network and network_interface.enabled:
dst_mac_address = self.software_manager.arp.get_arp_cache_mac_address(dst_ip_address)
break
if dst_mac_address:
use_default_gateway = False
outbound_network_interface = self.software_manager.arp.get_arp_cache_network_interface(dst_ip_address)
if use_default_gateway:
dst_mac_address = self.software_manager.arp.get_default_gateway_mac_address()
outbound_network_interface = self.software_manager.arp.get_default_gateway_network_interface()
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
ip_protocol: IPProtocol = IPProtocol.TCP,
icmp_packet: Optional[ICMPPacket] = None,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission.
@@ -157,74 +276,82 @@ class SessionManager:
: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 = session.with_ip_address
dst_port = session.dst_port
# Determine if the payload is for broadcast or unicast
# 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:
# 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
if isinstance(payload, ARPPacket):
# ARP requests need to be handles differently
if payload.request:
dst_mac_address = "ff:ff:ff:ff:ff:ff"
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
dst_mac_address = payload.target_mac_addr
outbound_network_interface = self.resolve_outbound_network_interface(payload.target_ip_address)
is_broadcast = payload.request
ip_protocol = IPProtocol.UDP
else:
vals = self.resolve_outbound_transmission_details(
dst_ip_address=dst_ip_address,
src_port=src_port,
dst_port=dst_port,
protocol=ip_protocol,
session_id=session_id,
)
(
outbound_network_interface,
dst_mac_address,
dst_ip_address,
src_port,
dst_port,
protocol,
is_broadcast,
) = vals
if protocol:
ip_protocol = protocol
# Check if outbound NIC and destination MAC address are resolved
if not outbound_nic or not dst_mac_address:
if not outbound_network_interface 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(
src_ip_address=outbound_nic.ip_address,
dst_ip_address=dst_ip_address,
),
tcp=TCPHeader(
if not (src_port or dst_port):
raise ValueError(
"Failed to resolve src or dst port. Have you sent the port from the service or application?"
)
tcp_header = None
udp_header = None
if ip_protocol == IPProtocol.TCP:
tcp_header = TCPHeader(
src_port=dst_port,
dst_port=dst_port,
)
elif ip_protocol == IPProtocol.UDP:
udp_header = UDPHeader(
src_port=dst_port,
dst_port=dst_port,
)
# TODO: Only create IP packet if not ARP
# ip_packet = None
# if dst_port != Port.ARP:
# IPPacket(
# src_ip_address=outbound_network_interface.ip_address,
# dst_ip_address=dst_ip_address,
# protocol=ip_protocol
# )
# Construct the frame for transmission
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_network_interface.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
src_ip_address=outbound_network_interface.ip_address,
dst_ip_address=dst_ip_address,
protocol=ip_protocol,
),
tcp=tcp_header,
udp=udp_header,
icmp=icmp_packet,
payload=payload,
)
# Manage session for unicast transmission
# TODO: Only create sessions for TCP
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)
@@ -235,9 +362,9 @@ class SessionManager:
self.sessions_by_uuid[session.uuid] = session
# Send the frame through the NIC
return outbound_nic.send_frame(frame)
return outbound_network_interface.send_frame(frame)
def receive_frame(self, frame: Frame):
def receive_frame(self, frame: Frame, from_network_interface: "NetworkInterface"):
"""
Receive a Frame.
@@ -246,6 +373,7 @@ class SessionManager:
:param frame: The frame being received.
"""
# TODO: Only create sessions for TCP
session_key = self._get_session_key(frame, inbound_frame=True)
session: Session = self.sessions_by_key.get(session_key)
if not session:
@@ -253,8 +381,20 @@ class SessionManager:
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
dst_port = None
if frame.tcp:
dst_port = frame.tcp.dst_port
elif frame.udp:
dst_port = frame.udp.dst_port
elif frame.icmp:
dst_port = Port.NONE
self.software_manager.receive_payload_from_session_manager(
payload=frame.payload, port=frame.tcp.dst_port, protocol=frame.ip.protocol, session_id=session.uuid
payload=frame.payload,
port=dst_port,
protocol=frame.ip.protocol,
session_id=session.uuid,
from_network_interface=from_network_interface,
frame=frame,
)
def show(self, markdown: bool = False):

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

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
@@ -14,7 +15,9 @@ from primaite.simulator.system.software import IOSoftware
if TYPE_CHECKING:
from primaite.simulator.system.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.base import Node, NIC
from primaite.simulator.system.services.arp.arp import ARP
from primaite.simulator.system.services.icmp.icmp import ICMP
from typing import Type, TypeVar
@@ -22,7 +25,14 @@ IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
class SoftwareManager:
"""A class that manages all running Services and Applications on a Node and facilitates their communication."""
"""
Manages all running services and applications on a network node and facilitates their communication.
This class is responsible for installing, uninstalling, and managing the operational state of various network
services and applications. It acts as a bridge between the node's session manager and its software components,
ensuring that incoming and outgoing network payloads are correctly routed to and from the appropriate services
or applications.
"""
def __init__(
self,
@@ -46,17 +56,26 @@ class SoftwareManager:
self.file_system: FileSystem = file_system
self.dns_server: Optional[IPv4Address] = dns_server
@property
def arp(self) -> "ARP":
"""Provides access to the ARP service instance, if installed."""
return self.software.get("ARP") # noqa
@property
def icmp(self) -> "ICMP":
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
def get_open_ports(self) -> List[Port]:
"""
Get a list of open ports.
:return: A list of all open ports on the Node.
"""
open_ports = [Port.ARP]
open_ports = []
for software in self.port_protocol_mapping.values():
if software.operating_state in {ApplicationOperatingState.RUNNING, ServiceOperatingState.RUNNING}:
open_ports.append(software.port)
open_ports.sort(key=lambda port: port.value)
return open_ports
def install(self, software_class: Type[IOSoftwareClass]):
@@ -132,6 +151,7 @@ class SoftwareManager:
payload: Any,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
session_id: Optional[str] = None,
) -> bool:
"""
@@ -151,10 +171,19 @@ class SoftwareManager:
payload=payload,
dst_ip_address=dest_ip_address,
dst_port=dest_port,
ip_protocol=ip_protocol,
session_id=session_id,
)
def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str):
def receive_payload_from_session_manager(
self,
payload: Any,
port: Port,
protocol: IPProtocol,
session_id: str,
from_network_interface: "NIC",
frame: Frame,
):
"""
Receive a payload from the SessionManager and forward it to the corresponding service or application.
@@ -163,7 +192,9 @@ class SoftwareManager:
"""
receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None)
if receiver:
receiver.receive(payload=payload, session_id=session_id)
receiver.receive(
payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame
)
else:
self.sys_log.error(f"No service or application found for port {port} and protocol {protocol}")
pass
@@ -174,7 +205,7 @@ class SoftwareManager:
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Name", "Type", "Operating State", "Health State", "Port"])
table = PrettyTable(["Name", "Type", "Operating State", "Health State", "Port", "Protocol"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -187,7 +218,8 @@ class SoftwareManager:
software_type,
software.operating_state.name,
software.health_state_actual.name,
software.port.value,
software.port.value if software.port != Port.NONE else None,
software.protocol.value,
]
)
print(table)

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

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

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

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

@@ -8,6 +8,7 @@ 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
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.core.sys_log import SysLog
@@ -229,6 +230,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."
protocol: IPProtocol
"The IP Protocol the Software operates on."
_connections: Dict[str, Dict] = {}
"Active connections."
@@ -334,6 +337,7 @@ class IOSoftware(Software):
session_id: Optional[str] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
**kwargs,
) -> bool:
"""
@@ -353,7 +357,11 @@ class IOSoftware(Software):
return False
return self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
payload=payload,
dest_ip_address=dest_ip_address,
dest_port=dest_port,
ip_protocol=ip_protocol,
session_id=session_id,
)
@abstractmethod

38
src/primaite/utils/validators.py Normal file
View File

@@ -0,0 +1,38 @@
from ipaddress import IPv4Address
from typing import Any, Final
from pydantic import BeforeValidator
from typing_extensions import Annotated
def ipv4_validator(v: Any) -> IPv4Address:
"""
Validate the input and ensure it can be converted to an IPv4Address instance.
This function takes an input `v`, and if it's not already an instance of IPv4Address, it tries to convert it to one.
If the conversion is successful, the IPv4Address instance is returned. This is useful for ensuring that any input
data is strictly in the format of an IPv4 address.
:param v: The input value that needs to be validated or converted to IPv4Address.
:return: An instance of IPv4Address.
:raises ValueError: If `v` is not a valid IPv4 address and cannot be converted to an instance of IPv4Address.
"""
if isinstance(v, IPv4Address):
return v
return IPv4Address(v)
# Define a custom type IPV4Address using the typing_extensions.Annotated.
# Annotated is used to attach metadata to type hints. In this case, it's used to associate the ipv4_validator
# with the IPv4Address type, ensuring that any usage of IPV4Address undergoes validation before assignment.
IPV4Address: Final[Annotated] = Annotated[IPv4Address, BeforeValidator(ipv4_validator)]
"""
IPv4Address with with IPv4Address with with pre-validation and auto-conversion from str using ipv4_validator..
This type is essentially an IPv4Address from the standard library's ipaddress module,
but with added validation logic. If you use this custom type, the ipv4_validator function
will automatically check and convert the input value to an instance of IPv4Address before
any Pydantic model uses it. This ensures that any field marked with this type is not just
an IPv4Address in form, but also valid according to the rules defined in ipv4_validator.
"""