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

Merge branch 'dev' into bugfix/2299-check_hash_function_corrupts_files_and_folders

Browse Source
This commit is contained in:
Nick Todd
2024-04-23 16:40:02 +01:00
parent 90224960e5 a60e8b2203
commit 3eea70f7d8
140 changed files with 9846 additions and 5182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -14,6 +14,7 @@ class _SimOutput:
)
self.save_pcap_logs: bool = False
self.save_sys_logs: bool = False
self.write_sys_log_to_terminal: bool = False
@property
def path(self) -> Path:

37
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Build a simulation using the Python API\n",
"\n",
"Currently, this notbook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n"
"Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n"
]
},
{
@@ -58,7 +58,8 @@
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.network.hardware.base import Node\n"
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server"
]
},
{
@@ -67,9 +68,9 @@
"metadata": {},
"outputs": [],
"source": [
"my_pc = Node(hostname=\"primaite_pc\",)\n",
"my_pc = Computer(hostname=\"Computer\", ip_address=\"192.168.1.10\", subnet_mask=\"255.255.255.0\")\n",
"net.add_node(my_pc)\n",
"my_server = Node(hostname=\"google_server\")\n",
"my_server = Server(hostname=\"Server\", ip_address=\"192.168.1.11\", subnet_mask=\"255.255.255.0\")\n",
"net.add_node(my_server)\n"
]
},
@@ -86,7 +87,8 @@
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.network.hardware.base import NIC, Link, Switch\n"
"from primaite.simulator.network.hardware.nodes.host.host_node import NIC\n",
"from primaite.simulator.network.hardware.nodes.network.switch import Switch\n"
]
},
{
@@ -95,19 +97,17 @@
"metadata": {},
"outputs": [],
"source": [
"my_swtich = Switch(hostname=\"switch1\", num_ports=12)\n",
"net.add_node(my_swtich)\n",
"my_switch = Switch(hostname=\"switch1\", num_ports=12)\n",
"net.add_node(my_switch)\n",
"\n",
"pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
"my_pc.connect_nic(pc_nic)\n",
"\n",
"\n",
"server_nic = NIC(ip_address=\"130.1.1.2\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
"my_server.connect_nic(server_nic)\n",
"\n",
"\n",
"net.connect(pc_nic, my_swtich.switch_ports[1])\n",
"net.connect(server_nic, my_swtich.switch_ports[2])\n"
"net.connect(pc_nic, my_switch.network_interface[1])\n",
"net.connect(server_nic, my_switch.network_interface[2])\n"
]
},
{
@@ -124,7 +124,8 @@
"outputs": [],
"source": [
"from primaite.simulator.file_system.file_type import FileType\n",
"from primaite.simulator.file_system.file_system import File"
"from primaite.simulator.file_system.file_system import File\n",
"from primaite.simulator.system.core.sys_log import SysLog"
]
},
{
@@ -134,7 +135,7 @@
"outputs": [],
"source": [
"my_pc_downloads_folder = my_pc.file_system.create_folder(\"downloads\")\n",
"my_pc_downloads_folder.add_file(File(name=\"firefox_installer.zip\",file_type=FileType.ZIP))"
"my_pc_downloads_folder.add_file(File(name=\"firefox_installer.zip\",folder_id=\"Test\", folder_name=\"downloads\" ,file_type=FileType.ZIP, sys_log=SysLog(hostname=\"Test\")))"
]
},
{
@@ -160,9 +161,12 @@
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"from primaite.simulator.system.applications.application import Application, ApplicationOperatingState\n",
"from primaite.simulator.system.software import SoftwareHealthState, SoftwareCriticality\n",
"from primaite.simulator.network.transmission.transport_layer import Port\n",
"from primaite.simulator.network.transmission.network_layer import IPProtocol\n",
"from primaite.simulator.file_system.file_system import FileSystem\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application):\n",
@@ -176,7 +180,7 @@
"metadata": {},
"outputs": [],
"source": [
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, ports={Port.HTTP}, operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual')"
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=Port.HTTP, protocol = IPProtocol.NONE,operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)"
]
},
{
@@ -257,9 +261,8 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
},
"orig_nbformat": 4
"version": "3.10.11"
}
},
"nbformat": 4,
"nbformat_minor": 2

148
src/primaite/simulator/_package_data/network_simulator_demo.ipynb
View File

@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "markdown",
"id": "03b2013a-b7d1-47ee-b08c-8dab83833720",
"id": "0",
"metadata": {},
"source": [
"# PrimAITE Router Simulation Demo\n",
@@ -12,7 +12,7 @@
},
{
"cell_type": "raw",
"id": "c8bb5698-e746-4e90-9c2f-efe962acdfa0",
"id": "1",
"metadata": {},
"source": [
" +------------+\n",
@@ -48,7 +48,7 @@
},
{
"cell_type": "markdown",
"id": "415d487c-6457-497d-85d6-99439b3541e7",
"id": "2",
"metadata": {},
"source": [
"## The Network\n",
@@ -60,7 +60,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "de57ac8c-5b28-4847-a759-2ceaf5593329",
"id": "3",
"metadata": {
"tags": []
},
@@ -72,7 +72,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "a1e2e4df-67c0-4584-ab27-47e2c7c7fcd2",
"id": "4",
"metadata": {
"tags": []
},
@@ -83,7 +83,7 @@
},
{
"cell_type": "markdown",
"id": "fb052c56-e9ca-4093-9115-d0c440b5ff53",
"id": "5",
"metadata": {},
"source": [
"Most of the Network components have a `.show()` function that prints a table of information about that object. We can view the Nodes and Links on the Network by calling `network.show()`."
@@ -92,7 +92,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "cc199741-ef2e-47f5-b2f0-e20049ccf40f",
"id": "6",
"metadata": {
"tags": []
},
@@ -103,7 +103,7 @@
},
{
"cell_type": "markdown",
"id": "76d2b7e9-280b-4741-a8b3-a84bed219fac",
"id": "7",
"metadata": {
"tags": []
},
@@ -115,7 +115,7 @@
},
{
"cell_type": "markdown",
"id": "84113002-843e-4cab-b899-667b50f25f6b",
"id": "8",
"metadata": {},
"source": [
"### Router Nodes\n",
@@ -125,7 +125,7 @@
},
{
"cell_type": "markdown",
"id": "bf63a178-eee5-4669-bf64-13aea7ecf6cb",
"id": "9",
"metadata": {},
"source": [
"Calling `router.show()` displays the Ethernet interfaces on the Router. If you need a table in markdown format, pass `markdown=True`."
@@ -134,7 +134,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "e76d1854-961e-438c-b40f-77fd9c3abe38",
"id": "10",
"metadata": {
"tags": []
},
@@ -145,7 +145,7 @@
},
{
"cell_type": "markdown",
"id": "e000540c-687c-4254-870c-1d814603bdbf",
"id": "11",
"metadata": {},
"source": [
"Calling `router.arp.show()` displays the Router ARP Cache."
@@ -154,7 +154,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "92de8b42-92d7-4934-9c12-50bf724c9eb2",
"id": "12",
"metadata": {
"tags": []
},
@@ -165,7 +165,7 @@
},
{
"cell_type": "markdown",
"id": "a9ff7ee8-9482-44de-9039-b684866bdc82",
"id": "13",
"metadata": {},
"source": [
"Calling `router.acl.show()` displays the Access Control List."
@@ -174,7 +174,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "5922282a-d22b-4e55-9176-f3f3654c849f",
"id": "14",
"metadata": {
"tags": []
},
@@ -185,7 +185,7 @@
},
{
"cell_type": "markdown",
"id": "71c87884-f793-4c9f-b004-5b0df86cf585",
"id": "15",
"metadata": {},
"source": [
"Calling `router.router_table.show()` displays the static routes the Router provides."
@@ -194,7 +194,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "327203be-f475-4727-82a1-e992d3b70ed8",
"id": "16",
"metadata": {
"tags": []
},
@@ -205,7 +205,7 @@
},
{
"cell_type": "markdown",
"id": "eef561a8-3d39-4c8b-bbc8-e8b10b8ed25f",
"id": "17",
"metadata": {},
"source": [
"Calling `router.sys_log.show()` displays the Router system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=<number of log entries>`."
@@ -214,7 +214,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "3d0aa004-b10c-445f-aaab-340e0e716c74",
"id": "18",
"metadata": {
"tags": []
},
@@ -225,7 +225,7 @@
},
{
"cell_type": "markdown",
"id": "25630c90-c54e-4b5d-8bf4-ad1b0722e126",
"id": "19",
"metadata": {},
"source": [
"### Switch Nodes\n",
@@ -235,16 +235,16 @@
},
{
"cell_type": "markdown",
"id": "4879394d-2981-40de-a229-e19b09a34e6e",
"id": "20",
"metadata": {},
"source": [
"Calling `switch.show()` displays the Switch orts on the Switch."
"Calling `switch.show()` displays the Switch ports on the Switch."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7fd439b-5442-4e9d-9e7d-86dacb77f458",
"id": "21",
"metadata": {
"tags": []
},
@@ -255,29 +255,7 @@
},
{
"cell_type": "markdown",
"id": "beb8dbd6-7250-4ac9-9fa2-d2a9c0e5fd19",
"metadata": {
"tags": []
},
"source": [
"Calling `switch.arp.show()` displays the Switch ARP Cache."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d06e1310-4a77-4315-a59f-cb1b49ca2352",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"switch_1\").arp.show()"
]
},
{
"cell_type": "markdown",
"id": "fda75ac3-8123-4234-8f36-86547891d8df",
"id": "22",
"metadata": {},
"source": [
"Calling `switch.sys_log.show()` displays the Switch system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=<number of log entries>`."
@@ -286,7 +264,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "a0d984b7-a7c1-4bbd-aa5a-9d3caecb08dc",
"id": "23",
"metadata": {
"tags": []
},
@@ -297,7 +275,7 @@
},
{
"cell_type": "markdown",
"id": "2f1d99ad-db4f-4baf-8a35-e1d95f269586",
"id": "24",
"metadata": {},
"source": [
"### Computer/Server Nodes\n",
@@ -307,7 +285,7 @@
},
{
"cell_type": "markdown",
"id": "c9e2251a-1b47-46e5-840f-7fec3e39c5aa",
"id": "25",
"metadata": {
"tags": []
},
@@ -318,7 +296,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "656c37f6-b145-42af-9714-8d2886d0eff8",
"id": "26",
"metadata": {
"tags": []
},
@@ -329,7 +307,7 @@
},
{
"cell_type": "markdown",
"id": "f1097a49-a3da-4d79-a06d-ae8af452918f",
"id": "27",
"metadata": {},
"source": [
"Calling `computer.arp.show()` displays the Computer/Server ARP Cache."
@@ -338,7 +316,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "66b267d6-2308-486a-b9aa-cb8d3bcf0753",
"id": "28",
"metadata": {
"tags": []
},
@@ -349,16 +327,16 @@
},
{
"cell_type": "markdown",
"id": "0d1fcad8-5b1a-4d8b-a49f-aa54a95fcaf0",
"id": "29",
"metadata": {},
"source": [
"Calling `switch.sys_log.show()` displays the Computer/Server system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=<number of log entries>`."
"Calling `computer.sys_log.show()` displays the Computer/Server system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=<number of log entries>`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1b5debe8-ef1b-445d-8fa9-6a45568f21f3",
"id": "30",
"metadata": {
"tags": []
},
@@ -369,7 +347,7 @@
},
{
"cell_type": "markdown",
"id": "fcfa1773-798c-4ada-9318-c3ad928217da",
"id": "31",
"metadata": {},
"source": [
"## Basic Network Comms Check\n",
@@ -380,7 +358,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "495b7de4-b6ce-41a6-9114-f74752ab4491",
"id": "32",
"metadata": {
"tags": []
},
@@ -391,7 +369,7 @@
},
{
"cell_type": "markdown",
"id": "3e13922a-217f-4f4e-99b6-57a07613cade",
"id": "33",
"metadata": {},
"source": [
"We'll first ping client_1's default gateway."
@@ -400,7 +378,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "a38abb71-994e-49e8-8f51-e9a550e95b99",
"id": "34",
"metadata": {
"tags": []
},
@@ -412,7 +390,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "8388e1e9-30e3-4534-8e5a-c6e9144149d2",
"id": "35",
"metadata": {
"tags": []
},
@@ -423,7 +401,7 @@
},
{
"cell_type": "markdown",
"id": "02c76d5c-d954-49db-912d-cb9c52f46375",
"id": "36",
"metadata": {},
"source": [
"Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)."
@@ -432,7 +410,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "ff8e976a-c16b-470c-8923-325713a30d6c",
"id": "37",
"metadata": {
"tags": []
},
@@ -443,7 +421,7 @@
},
{
"cell_type": "markdown",
"id": "80280404-a5ab-452f-8a02-771a0d7496b1",
"id": "38",
"metadata": {},
"source": [
"And finally, we'll ping the web server."
@@ -452,7 +430,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "c4163f8d-6a72-410c-9f5c-4f881b7de45e",
"id": "39",
"metadata": {
"tags": []
},
@@ -463,7 +441,7 @@
},
{
"cell_type": "markdown",
"id": "1194c045-ba77-4427-be30-ed7b5b224850",
"id": "40",
"metadata": {},
"source": [
"To confirm that the ping was received and processed by the web_server, we can view the sys log"
@@ -472,7 +450,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "e79a523a-5780-45b6-8798-c434e0e522bd",
"id": "41",
"metadata": {
"tags": []
},
@@ -483,17 +461,17 @@
},
{
"cell_type": "markdown",
"id": "5928f6dd-1006-45e3-99f3-8f311a875faa",
"id": "42",
"metadata": {},
"source": [
"## Advanced Network Usage\n",
"\n",
"We can now use the Network to perform some more advaced things."
"We can now use the Network to perform some more advanced things."
]
},
{
"cell_type": "markdown",
"id": "5e023ef3-7d18-4006-96ee-042a06a481fc",
"id": "43",
"metadata": {},
"source": [
"Let's attempt to prevent client_2 from being able to ping the web server. First, we'll confirm that it can ping the server first..."
@@ -502,7 +480,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "603cf913-e261-49da-a7dd-85e1bb6dec56",
"id": "44",
"metadata": {
"tags": []
},
@@ -513,7 +491,7 @@
},
{
"cell_type": "markdown",
"id": "5cf962a4-20e6-44ae-9748-7fc5267ae111",
"id": "45",
"metadata": {},
"source": [
"If we look at the client_2 sys log we can see that the four ICMP echo requests were sent and four ICMP each replies were received:"
@@ -522,7 +500,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "e047de00-3de4-4823-b26a-2c8d64c7a663",
"id": "46",
"metadata": {
"tags": []
},
@@ -533,7 +511,7 @@
},
{
"cell_type": "markdown",
"id": "bdc4741d-6e3e-4aec-a69c-c2e9653bd02c",
"id": "47",
"metadata": {},
"source": [
"Now we'll add an ACL to block ICMP from 192.168.10.22"
@@ -542,7 +520,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "6db355ae-b99a-441b-a2c4-4ffe78f46bff",
"id": "48",
"metadata": {
"tags": []
},
@@ -550,7 +528,7 @@
"source": [
"from primaite.simulator.network.transmission.network_layer import IPProtocol\n",
"from primaite.simulator.network.transmission.transport_layer import Port\n",
"from primaite.simulator.network.hardware.nodes.router import ACLAction\n",
"from primaite.simulator.network.hardware.nodes.network.router import ACLAction\n",
"network.get_node_by_hostname(\"router_1\").acl.add_rule(\n",
" action=ACLAction.DENY,\n",
" protocol=IPProtocol.ICMP,\n",
@@ -562,7 +540,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "a345e000-8842-4827-af96-adc0fbe390fb",
"id": "49",
"metadata": {
"tags": []
},
@@ -573,7 +551,7 @@
},
{
"cell_type": "markdown",
"id": "3a5bfd9f-04cb-493e-a86c-cd268563a262",
"id": "50",
"metadata": {},
"source": [
"Now we attempt (and fail) to ping the web server"
@@ -582,7 +560,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "a4f4ff31-590f-40fb-b13d-efaa8c2720b6",
"id": "51",
"metadata": {
"tags": []
},
@@ -593,7 +571,7 @@
},
{
"cell_type": "markdown",
"id": "83e56497-097b-45cb-964e-b15c72547b38",
"id": "52",
"metadata": {},
"source": [
"We can check that the ping was actually sent by client_2 by viewing the sys log"
@@ -602,7 +580,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "f62b8a4e-fd3b-4059-b108-3d4a0b18f2a0",
"id": "53",
"metadata": {
"tags": []
},
@@ -613,7 +591,7 @@
},
{
"cell_type": "markdown",
"id": "c7040311-a879-4620-86a0-55d0774156e5",
"id": "54",
"metadata": {},
"source": [
"We can check the router sys log to see why the traffic was blocked"
@@ -622,7 +600,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "7e53d776-99da-4d2c-a2a7-bd7ce27bff4c",
"id": "55",
"metadata": {
"tags": []
},
@@ -633,7 +611,7 @@
},
{
"cell_type": "markdown",
"id": "aba0bc7d-da57-477b-b34a-3688b5aab2c6",
"id": "56",
"metadata": {},
"source": [
"Now a final check to ensure that client_1 can still ping the web_server."
@@ -642,7 +620,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "d542734b-7582-4af7-8254-bda3de50d091",
"id": "57",
"metadata": {
"tags": []
},
@@ -654,7 +632,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "d78e9fe3-02c6-4792-944f-5622e26e0412",
"id": "58",
"metadata": {
"tags": []
},

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

@@ -7,12 +7,10 @@ from uuid import uuid4
from pydantic import BaseModel, ConfigDict, Field, validate_call
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.interface.request import RequestFormat, RequestResponse
_LOGGER = getLogger(__name__)
RequestFormat = List[Union[str, int, float]]
class RequestPermissionValidator(BaseModel):
"""
@@ -228,6 +226,15 @@ class SimComponent(BaseModel):
return
return self._request_manager(request, context)
def pre_timestep(self, timestep: int) -> None:
"""
Apply any logic that needs to happen at the beginning of the timestep to ensure correct observations/rewards.
:param timestep: what's the current time
:type timestep: int
"""
pass
def apply_timestep(self, timestep: int) -> None:
"""
Apply a timestep evolution to this component.

4
src/primaite/simulator/file_system/file.py
View File

@@ -103,6 +103,10 @@ class File(FileSystemItemABC):
"""
super().apply_timestep(timestep=timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
# reset the number of accesses to 0
self.num_access = 0

12
src/primaite/simulator/file_system/file_system.py
View File

@@ -427,15 +427,21 @@ class FileSystem(SimComponent):
"""Apply time step to FileSystem and its child folders and files."""
super().apply_timestep(timestep=timestep)
# apply timestep to folders
for folder_id in self.folders:
self.folders[folder_id].apply_timestep(timestep=timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
# reset number of file creations
self.num_file_creations = 0
# reset number of file deletions
self.num_file_deletions = 0
# apply timestep to folders
for folder_id in self.folders:
self.folders[folder_id].apply_timestep(timestep=timestep)
for folder in self.folders.values():
folder.pre_timestep(timestep)
###############################################################
# Agent actions

7
src/primaite/simulator/file_system/folder.py
View File

@@ -128,6 +128,13 @@ class Folder(FileSystemItemABC):
for file_id in self.files:
self.files[file_id].apply_timestep(timestep=timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
for file in self.files.values():
file.pre_timestep(timestep)
def _scan_timestep(self) -> None:
"""Apply the scan action timestep."""
if self.scan_countdown >= 0:

4
src/primaite/simulator/network/airspace.py
View File

@@ -157,7 +157,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
return
if not self._connected_node:
_LOGGER.error(f"Interface {self} cannot be enabled as it is not connected to a Node")
_LOGGER.warning(f"Interface {self} cannot be enabled as it is not connected to a Node")
return
if self._connected_node.operating_state != NodeOperatingState.ON:
@@ -271,7 +271,7 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
state["frequency"] = self.frequency
state["frequency"] = self.frequency.value
return state

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

@@ -1,3 +1,4 @@
from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional
import matplotlib.pyplot as plt
@@ -8,6 +9,7 @@ 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, Node, WiredNetworkInterface
from primaite.simulator.network.hardware.nodes.host.server import Printer
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service
@@ -85,6 +87,16 @@ class Network(SimComponent):
for link_id in self.links:
self.links[link_id].apply_timestep(timestep=timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
for node in self.nodes.values():
node.pre_timestep(timestep)
for link in self.links.values():
link.pre_timestep(timestep)
@property
def router_nodes(self) -> List[Node]:
"""The Routers in the Network."""
@@ -110,6 +122,16 @@ class Network(SimComponent):
"""The Firewalls in the Network."""
return [node for node in self.nodes.values() if node.__class__.__name__ == "Firewall"]
@property
def printer_nodes(self) -> List[Node]:
"""The printers on the network."""
return [node for node in self.nodes.values() if isinstance(node, Printer)]
@property
def wireless_router_nodes(self) -> List[Node]:
"""The Routers in the Network."""
return [node for node in self.nodes.values() if node.__class__.__name__ == "WirelessRouter"]
def show(self, nodes: bool = True, ip_addresses: bool = True, links: bool = True, markdown: bool = False):
"""
Print tables describing the Network.
@@ -128,6 +150,8 @@ class Network(SimComponent):
"Switch": self.switch_nodes,
"Server": self.server_nodes,
"Computer": self.computer_nodes,
"Printer": self.printer_nodes,
"Wireless Router": self.wireless_router_nodes,
}
if nodes:
table = PrettyTable(["Node", "Type", "Operating State"])
@@ -150,14 +174,17 @@ class Network(SimComponent):
for node in nodes:
for i, port in node.network_interface.items():
if hasattr(port, "ip_address"):
port_str = port.port_name if port.port_name else port.port_num
table.add_row(
[node.hostname, port_str, port.ip_address, port.subnet_mask, node.default_gateway]
)
if port.ip_address != IPv4Address("127.0.0.1"):
port_str = port.port_name if port.port_name else port.port_num
table.add_row(
[node.hostname, port_str, port.ip_address, port.subnet_mask, node.default_gateway]
)
print(table)
if links:
table = PrettyTable(["Endpoint A", "Endpoint B", "is Up", "Bandwidth (MBits)", "Current Load"])
table = PrettyTable(
["Endpoint A", "A Port", "Endpoint B", "B Port", "is Up", "Bandwidth (MBits)", "Current Load"]
)
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -170,7 +197,9 @@ class Network(SimComponent):
table.add_row(
[
link.endpoint_a.parent.hostname,
str(link.endpoint_a),
link.endpoint_b.parent.hostname,
str(link.endpoint_b),
link.is_up,
link.bandwidth,
link.current_load_percent,
@@ -208,18 +237,19 @@ class Network(SimComponent):
}
)
# Update the links one-by-one. The key is a 4-tuple of `hostname_a, port_a, hostname_b, port_b`
for uuid, link in self.links.items():
for _, link in self.links.items():
node_a = link.endpoint_a._connected_node
node_b = link.endpoint_b._connected_node
hostname_a = node_a.hostname if node_a else None
hostname_b = node_b.hostname if node_b else None
port_a = link.endpoint_a.port_num
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
state["links"][uuid]["port_a"] = port_a
state["links"][uuid]["port_b"] = port_b
link_key = f"{hostname_a}:eth-{port_a}<->{hostname_b}:eth-{port_b}"
state["links"][link_key] = link.describe_state()
state["links"][link_key]["hostname_a"] = hostname_a
state["links"][link_key]["hostname_b"] = hostname_b
state["links"][link_key]["port_a"] = port_a
state["links"][link_key]["port_b"] = port_b
return state

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

@@ -9,7 +9,7 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
def num_of_switches_required(num_nodes: int, max_switch_ports: int = 24) -> int:
def num_of_switches_required(num_nodes: int, max_network_interface: int = 24) -> int:
"""
Calculate the minimum number of network switches required to connect a given number of nodes.
@@ -18,7 +18,7 @@ def num_of_switches_required(num_nodes: int, max_switch_ports: int = 24) -> int:
to accommodate all nodes under this constraint.
:param num_nodes: The total number of nodes that need to be connected in the network.
:param max_switch_ports: The maximum number of ports available on each switch. Defaults to 24.
:param max_network_interface: The maximum number of ports available on each switch. Defaults to 24.
:return: The minimum number of switches required to connect all PCs.
@@ -33,11 +33,11 @@ def num_of_switches_required(num_nodes: int, max_switch_ports: int = 24) -> int:
3
"""
# Reduce the effective number of switch ports by 1 to leave space for the router
effective_switch_ports = max_switch_ports - 1
effective_network_interface = max_network_interface - 1
# Calculate the number of fully utilised switches and any additional switch for remaining PCs
full_switches = num_nodes // effective_switch_ports
extra_pcs = num_nodes % effective_switch_ports
full_switches = num_nodes // effective_network_interface
extra_pcs = num_nodes % effective_network_interface
# Return the total number of switches required
return full_switches + (1 if extra_pcs > 0 else 0)
@@ -77,7 +77,7 @@ def create_office_lan(
# Calculate the required number of switches
num_of_switches = num_of_switches_required(num_nodes=num_pcs)
effective_switch_ports = 23 # One port less for router connection
effective_network_interface = 23 # One port less for router connection
if pcs_ip_block_start <= num_of_switches:
raise ValueError(f"pcs_ip_block_start must be greater than the number of required switches {num_of_switches}")
@@ -116,7 +116,7 @@ def create_office_lan(
# Add PCs to the LAN and connect them to switches
for i in range(1, num_pcs + 1):
# Add a new edge switch if the current one is full
if switch_port == effective_switch_ports:
if switch_port == effective_network_interface:
switch_n += 1
switch_port = 0
switch = Switch(hostname=f"switch_edge_{switch_n}_{lan_name}", start_up_duration=0)

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

@@ -5,7 +5,7 @@ import secrets
from abc import ABC, abstractmethod
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Optional, Type, TypeVar, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, Field
@@ -35,8 +35,11 @@ from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.processes.process import Process
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import IOSoftware
from primaite.utils.validators import IPV4Address
IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
_LOGGER = getLogger(__name__)
@@ -108,7 +111,7 @@ class NetworkInterface(SimComponent, ABC):
"""Reset the original state of the SimComponent."""
super().setup_for_episode(episode=episode)
self.nmne = {}
if episode and self.pcap:
if episode and self.pcap and SIM_OUTPUT.save_pcap_logs:
self.pcap.current_episode = episode
self.pcap.setup_logger()
self.enable()
@@ -261,6 +264,9 @@ class NetworkInterface(SimComponent, ABC):
"""
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}"
def __hash__(self) -> int:
return hash(self.uuid)
def apply_timestep(self, timestep: int) -> None:
"""
Apply a timestep evolution to this component.
@@ -297,7 +303,7 @@ class WiredNetworkInterface(NetworkInterface, ABC):
return True
if not self._connected_node:
_LOGGER.error(f"Interface {self} cannot be enabled as it is not connected to a Node")
_LOGGER.warning(f"Interface {self} cannot be enabled as it is not connected to a Node")
return False
if self._connected_node.operating_state != NodeOperatingState.ON:
@@ -343,11 +349,11 @@ class WiredNetworkInterface(NetworkInterface, ABC):
:param link: The Link instance to connect to this network interface.
"""
if self._connected_link:
_LOGGER.error(f"Cannot connect Link to network interface {self} as it already has a connection")
_LOGGER.warning(f"Cannot connect Link to network interface {self} as it already has a connection")
return
if self._connected_link == link:
_LOGGER.error(f"Cannot connect Link to network interface {self} as it is already connected")
_LOGGER.warning(f"Cannot connect Link to network interface {self} as it is already connected")
return
self._connected_link = link
@@ -519,12 +525,10 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC):
"""
super().enable()
try:
pass
self._connected_node.default_gateway_hello()
return True
except AttributeError:
pass
return False
return True
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
@@ -660,6 +664,10 @@ class Link(SimComponent):
def apply_timestep(self, timestep: int) -> None:
"""Apply a timestep to the simulation."""
super().apply_timestep(timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
self.current_load = 0.0
@@ -845,12 +853,62 @@ class Node(SimComponent):
)
rm.add_request("os", RequestType(func=self._os_request_manager, validator=_node_is_on))
self._software_request_manager = RequestManager()
rm.add_request("software_manager", RequestType(func=self._software_request_manager, validator=_node_is_on))
self._application_manager = RequestManager()
self._software_request_manager.add_request(
name="application", request_type=RequestType(func=self._application_manager)
)
self._application_manager.add_request(
name="install",
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(
self.application_install_action(
application=self._read_application_type(request[0]), ip_address=request[1]
)
)
),
)
self._application_manager.add_request(
name="uninstall",
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(
self.application_uninstall_action(application=self._read_application_type(request[0]))
)
),
)
return rm
def _install_system_software(self):
"""Install System Software - software that is usually provided with the OS."""
pass
def _read_application_type(self, application_class_str: str) -> Type[IOSoftwareClass]:
"""Wrapper that converts the string from the request manager into the appropriate class for the application."""
if application_class_str == "DoSBot":
from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
return DoSBot
elif application_class_str == "DataManipulationBot":
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import (
DataManipulationBot,
)
return DataManipulationBot
elif application_class_str == "WebBrowser":
from primaite.simulator.system.applications.web_browser import WebBrowser
return WebBrowser
elif application_class_str == "RansomwareScript":
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
return RansomwareScript
else:
return 0
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -891,8 +949,9 @@ class Node(SimComponent):
table.align = "l"
table.title = f"{self.hostname} Open Ports"
for port in self.software_manager.get_open_ports():
table.add_row([port.value, port.name])
print(table)
if port.value > 0:
table.add_row([port.value, port.name])
print(table.get_string(sortby="Port"))
@property
def has_enabled_network_interface(self) -> bool:
@@ -917,12 +976,15 @@ class Node(SimComponent):
table.align = "l"
table.title = f"{self.hostname} Network Interface Cards"
for port, network_interface in self.network_interface.items():
ip_address = ""
if hasattr(network_interface, "ip_address"):
ip_address = f"{network_interface.ip_address}/{network_interface.ip_network.prefixlen}"
table.add_row(
[
port,
type(network_interface),
network_interface.__class__.__name__,
network_interface.mac_address,
f"{network_interface.ip_address}/{network_interface.ip_network.prefixlen}",
ip_address,
network_interface.speed,
"Enabled" if network_interface.enabled else "Disabled",
]
@@ -1023,6 +1085,23 @@ class Node(SimComponent):
self.file_system.apply_timestep(timestep=timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
for network_interface in self.network_interfaces.values():
network_interface.pre_timestep(timestep=timestep)
for process_id in self.processes:
self.processes[process_id].pre_timestep(timestep=timestep)
for service_id in self.services:
self.services[service_id].pre_timestep(timestep=timestep)
for application_id in self.applications:
self.applications[application_id].pre_timestep(timestep=timestep)
self.file_system.pre_timestep(timestep=timestep)
def scan(self) -> bool:
"""
Scan the node and all the items within it.
@@ -1259,6 +1338,77 @@ class Node(SimComponent):
_LOGGER.info(f"Removed application {application.name} from node {self.hostname}")
self._application_request_manager.remove_request(application.name)
def application_install_action(self, application: Application, ip_address: Optional[str] = None) -> bool:
"""
Install an application on this node and configure it.
This method is useful for allowing agents to take this action.
:param application: Application object that has not been installed on any node yet.
:type application: Application
:param ip_address: IP address used to configure the application
(target IP for the DoSBot or server IP for the DataManipulationBot)
:type ip_address: str
:return: True if the application is installed successfully, otherwise False.
"""
if application in self:
_LOGGER.warning(
f"Can't add application {application.__name__}" + f"to node {self.hostname}. It's already installed."
)
return True
self.software_manager.install(application)
application_instance = self.software_manager.software.get(str(application.__name__))
self.applications[application_instance.uuid] = application_instance
self.sys_log.info(f"Installed application {application_instance.name}")
_LOGGER.debug(f"Added application {application_instance.name} to node {self.hostname}")
self._application_request_manager.add_request(
application_instance.name, RequestType(func=application_instance._request_manager)
)
# Configure application if additional parameters are given
if ip_address:
if application_instance.name == "DoSBot":
application_instance.configure(target_ip_address=IPv4Address(ip_address))
elif application_instance.name == "DataManipulationBot":
application_instance.configure(server_ip_address=IPv4Address(ip_address))
elif application_instance.name == "RansomwareScript":
application_instance.configure(server_ip_address=IPv4Address(ip_address))
else:
pass
if application_instance.name in self.software_manager.software:
return True
else:
return False
def application_uninstall_action(self, application: Application) -> bool:
"""
Uninstall and completely remove application from this node.
This method is useful for allowing agents to take this action.
:param application: Application object that is currently associated with this node.
:type application: Application
:return: True if the application is uninstalled successfully, otherwise False.
"""
if application.__name__ not in self.software_manager.software:
_LOGGER.warning(
f"Can't remove application {application.__name__}" + f"from node {self.hostname}. It's not installed."
)
return True
application_instance = self.software_manager.software.get(
str(application.__name__)
) # This works because we can't have two applications with the same name on the same node
# self.uninstall_application(application_instance)
self.software_manager.uninstall(application_instance.name)
if application_instance.name not in self.software_manager.software:
return True
else:
return False
def _shut_down_actions(self):
"""Actions to perform when the node is shut down."""
# Turn off all the services in the node
@@ -1290,4 +1440,6 @@ class Node(SimComponent):
def __contains__(self, item: Any) -> bool:
if isinstance(item, Service):
return item.uuid in self.services
elif isinstance(item, Application):
return item.uuid in self.applications
return None

10
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -316,6 +316,16 @@ class HostNode(Node):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
@property
def arp(self) -> Optional[ARP]:
"""
Return the ARP Cache of the HostNode.
:return: ARP Cache for given HostNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
def _install_system_software(self):
"""
Installs the system software and network services typically found on an operating system.

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

@@ -28,3 +28,9 @@ class Server(HostNode):
* Applications:
* Web Browser
"""
class Printer(HostNode):
"""Printer? I don't even know her!."""
# TODO: Implement printer-specific behaviour

129
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -1,9 +1,10 @@
from ipaddress import IPv4Address
from typing import Dict, Final, Optional, Union
from typing import Dict, Final, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from pydantic import Field, validate_call
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.network.router import (
AccessControlList,
@@ -67,22 +68,34 @@ class Firewall(Router):
:ivar str hostname: The Firewall hostname.
"""
internal_inbound_acl: Optional[AccessControlList] = None
internal_inbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="Internal Inbound", implicit_action=ACLAction.DENY)
)
"""Access Control List for managing entering the internal network."""
internal_outbound_acl: Optional[AccessControlList] = None
internal_outbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="Internal Outbound", implicit_action=ACLAction.DENY)
)
"""Access Control List for managing traffic leaving the internal network."""
dmz_inbound_acl: Optional[AccessControlList] = None
dmz_inbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="DMZ Inbound", implicit_action=ACLAction.DENY)
)
"""Access Control List for managing traffic entering the DMZ."""
dmz_outbound_acl: Optional[AccessControlList] = None
dmz_outbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="DMZ Outbound", implicit_action=ACLAction.DENY)
)
"""Access Control List for managing traffic leaving the DMZ."""
external_inbound_acl: Optional[AccessControlList] = None
external_inbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="External Inbound", implicit_action=ACLAction.PERMIT)
)
"""Access Control List for managing traffic entering from an external network."""
external_outbound_acl: Optional[AccessControlList] = None
external_outbound_acl: AccessControlList = Field(
default_factory=lambda: AccessControlList(name="External Outbound", implicit_action=ACLAction.PERMIT)
)
"""Access Control List for managing traffic leaving towards an external network."""
def __init__(self, hostname: str, **kwargs):
@@ -100,29 +113,85 @@ class Firewall(Router):
self.connect_nic(
RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", port_name="dmz")
)
# Update ACL objects with firewall's hostname and syslog to allow accurate logging
self.internal_inbound_acl.sys_log = kwargs["sys_log"]
self.internal_inbound_acl.name = f"{hostname} - Internal Inbound"
# 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.sys_log = kwargs["sys_log"]
self.internal_outbound_acl.name = f"{hostname} - Internal Outbound"
self.dmz_inbound_acl.sys_log = kwargs["sys_log"]
self.dmz_inbound_acl.name = f"{hostname} - DMZ Inbound"
self.dmz_outbound_acl.sys_log = kwargs["sys_log"]
self.dmz_outbound_acl.name = f"{hostname} - DMZ Outbound"
self.external_inbound_acl.sys_log = kwargs["sys_log"]
self.external_inbound_acl.name = f"{hostname} - External Inbound"
self.external_outbound_acl.sys_log = kwargs["sys_log"]
self.external_outbound_acl.name = f"{hostname} - External Outbound"
def _init_request_manager(self) -> RequestManager:
"""
Initialise the request manager.
More information in user guide and docstring for SimComponent._init_request_manager.
"""
rm = super()._init_request_manager()
self._internal_acl_request_manager = RequestManager()
rm.add_request("internal", RequestType(func=self._internal_acl_request_manager))
self._dmz_acl_request_manager = RequestManager()
rm.add_request("dmz", RequestType(func=self._dmz_acl_request_manager))
self._external_acl_request_manager = RequestManager()
rm.add_request("external", RequestType(func=self._external_acl_request_manager))
self._internal_inbound_acl_request_manager = RequestManager()
self._internal_outbound_acl_request_manager = RequestManager()
self._internal_acl_request_manager.add_request(
"inbound", RequestType(func=self._internal_inbound_acl_request_manager)
)
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"
self._internal_acl_request_manager.add_request(
"outbound", RequestType(func=self._internal_outbound_acl_request_manager)
)
# 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.dmz_inbound_acl_request_manager = RequestManager()
self.dmz_outbound_acl_request_manager = RequestManager()
self._dmz_acl_request_manager.add_request("inbound", RequestType(func=self.dmz_inbound_acl_request_manager))
self._dmz_acl_request_manager.add_request("outbound", RequestType(func=self.dmz_outbound_acl_request_manager))
self.external_inbound_acl_request_manager = RequestManager()
self.external_outbound_acl_request_manager = RequestManager()
self._external_acl_request_manager.add_request(
"inbound", RequestType(func=self.external_inbound_acl_request_manager)
)
self.external_outbound_acl = AccessControlList(
sys_log=kwargs["sys_log"], implicit_action=ACLAction.PERMIT, name=f"{hostname} - External Outbound"
self._external_acl_request_manager.add_request(
"outbound", RequestType(func=self.external_outbound_acl_request_manager)
)
self._internal_inbound_acl_request_manager.add_request(
"acl", RequestType(func=self.internal_inbound_acl._request_manager)
)
self._internal_outbound_acl_request_manager.add_request(
"acl", RequestType(func=self.internal_outbound_acl._request_manager)
)
self.dmz_inbound_acl_request_manager.add_request("acl", RequestType(func=self.dmz_inbound_acl._request_manager))
self.dmz_outbound_acl_request_manager.add_request(
"acl", RequestType(func=self.dmz_outbound_acl._request_manager)
)
self.external_inbound_acl_request_manager.add_request(
"acl", RequestType(func=self.external_inbound_acl._request_manager)
)
self.external_outbound_acl_request_manager.add_request(
"acl", RequestType(func=self.external_outbound_acl._request_manager)
)
return rm
def describe_state(self) -> Dict:
"""
Describes the current state of the Firewall.
@@ -530,7 +599,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
@@ -543,7 +614,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
@@ -556,7 +629,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
@@ -569,7 +644,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
@@ -582,7 +659,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
@@ -595,7 +674,9 @@ class Firewall(Router):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)

12
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -1,7 +1,9 @@
from abc import abstractmethod
from typing import Optional
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node):
@@ -28,3 +30,13 @@ class NetworkNode(Node):
:type from_network_interface: NetworkInterface
"""
pass
@property
def arp(self) -> Optional[ARP]:
"""
Return the ARP Cache of the NetworkNode.
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")

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

@@ -18,6 +18,7 @@ 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.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.services.arp.arp import ARP
from primaite.simulator.system.services.icmp.icmp import ICMP
@@ -147,8 +148,10 @@ class ACLRule(SimComponent):
state["action"] = self.action.value
state["protocol"] = self.protocol.name if self.protocol else None
state["src_ip_address"] = str(self.src_ip_address) if self.src_ip_address else None
state["src_wildcard_mask"] = str(self.src_wildcard_mask) if self.src_wildcard_mask else None
state["src_port"] = self.src_port.name if self.src_port else None
state["dst_ip_address"] = str(self.dst_ip_address) if self.dst_ip_address else None
state["dst_wildcard_mask"] = str(self.dst_wildcard_mask) if self.dst_wildcard_mask else None
state["dst_port"] = self.dst_port.name if self.dst_port else None
state["match_count"] = self.match_count
return state
@@ -275,7 +278,7 @@ class AccessControlList(SimComponent):
:ivar int max_acl_rules: The maximum number of ACL rules that can be added to the list. Defaults to 25.
"""
sys_log: SysLog
sys_log: Optional[SysLog] = None
implicit_action: ACLAction
implicit_rule: ACLRule
max_acl_rules: int = 25
@@ -319,10 +322,12 @@ class AccessControlList(SimComponent):
action=ACLAction[request[0]],
protocol=None if request[1] == "ALL" else IPProtocol[request[1]],
src_ip_address=None if request[2] == "ALL" else IPv4Address(request[2]),
src_port=None if request[3] == "ALL" else Port[request[3]],
dst_ip_address=None if request[4] == "ALL" else IPv4Address(request[4]),
dst_port=None if request[5] == "ALL" else Port[request[5]],
position=int(request[6]),
src_wildcard_mask=None if request[3] == "NONE" else IPv4Address(request[3]),
src_port=None if request[4] == "ALL" else Port[request[4]],
dst_ip_address=None if request[5] == "ALL" else IPv4Address(request[5]),
dst_wildcard_mask=None if request[6] == "NONE" else IPv4Address(request[6]),
dst_port=None if request[7] == "ALL" else Port[request[7]],
position=int(request[8]),
)
)
),
@@ -624,11 +629,12 @@ class RouteTable(SimComponent):
"""
pass
@validate_call()
def add_route(
self,
address: Union[IPv4Address, str],
subnet_mask: Union[IPv4Address, str],
next_hop_ip_address: Union[IPv4Address, str],
address: Union[IPV4Address, str],
subnet_mask: Union[IPV4Address, str],
next_hop_ip_address: Union[IPV4Address, str],
metric: float = 0.0,
):
"""
@@ -647,7 +653,8 @@ class RouteTable(SimComponent):
)
self.routes.append(route)
def set_default_route_next_hop_ip_address(self, ip_address: IPv4Address):
@validate_call()
def set_default_route_next_hop_ip_address(self, ip_address: IPV4Address):
"""
Sets the next-hop IP address for the default route in a routing table.
@@ -660,7 +667,7 @@ class RouteTable(SimComponent):
"""
if not self.default_route:
self.default_route = RouteEntry(
ip_address=IPv4Address("0.0.0.0"),
address=IPv4Address("0.0.0.0"),
subnet_mask=IPv4Address("0.0.0.0"),
next_hop_ip_address=ip_address,
)
@@ -767,6 +774,13 @@ class RouterARP(ARP):
is_reattempt=True,
is_default_route_attempt=is_default_route_attempt,
)
elif route and route == self.router.route_table.default_route:
self.send_arp_request(self.router.route_table.default_route.next_hop_ip_address)
return self._get_arp_cache_mac_address(
ip_address=self.router.route_table.default_route.next_hop_ip_address,
is_reattempt=True,
is_default_route_attempt=True,
)
else:
if self.router.route_table.default_route:
if not is_default_route_attempt:
@@ -817,6 +831,12 @@ class RouterARP(ARP):
return network_interface
if not is_reattempt:
if self.router.ip_is_in_router_interface_subnet(ip_address):
self.send_arp_request(ip_address)
return self._get_arp_cache_network_interface(
ip_address=ip_address, is_reattempt=True, is_default_route_attempt=is_default_route_attempt
)
route = self.router.route_table.find_best_route(ip_address)
if route and route != self.router.route_table.default_route:
self.send_arp_request(route.next_hop_ip_address)
@@ -825,6 +845,13 @@ class RouterARP(ARP):
is_reattempt=True,
is_default_route_attempt=is_default_route_attempt,
)
elif route and route == self.router.route_table.default_route:
self.send_arp_request(self.router.route_table.default_route.next_hop_ip_address)
return self._get_arp_cache_network_interface(
ip_address=self.router.route_table.default_route.next_hop_ip_address,
is_reattempt=True,
is_default_route_attempt=True,
)
else:
if self.router.route_table.default_route:
if not is_default_route_attempt:
@@ -1016,6 +1043,144 @@ class RouterInterface(IPWiredNetworkInterface):
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
class RouterSessionManager(SessionManager):
"""
Manages network sessions, including session creation, lookup, and communication with other components.
The RouterSessionManager is a Router/Firewall specific implementation of SessionManager. It overrides the
resolve_outbound_network_interface and resolve_outbound_transmission_details functions, allowing them to leverage
the route table instead of the default gateway.
:param sys_log: A reference to the system log component.
"""
def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional[RouterInterface]:
"""
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 performing a route table look-up to determine if there is a dedicated route or a default
route it can use.
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 according to route table look-up.
: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[RouterInterface]
"""
network_interface = super().resolve_outbound_network_interface(dst_ip_address)
if not network_interface:
route = self.node.route_table.find_best_route(dst_ip_address)
if not route:
return None
network_interface = super().resolve_outbound_network_interface(route.next_hop_ip_address)
return 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[RouterInterface],
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 defaults to performing a route table look-up to determine if there is a dedicated route
or a default route it can use.
: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[RouterInterface], 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_route_table = 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_route_table = False
outbound_network_interface = self.software_manager.arp.get_arp_cache_network_interface(dst_ip_address)
if use_route_table:
route = self.node.route_table.find_best_route(dst_ip_address)
if not route:
raise Exception("cannot use route to resolve outbound details")
dst_mac_address = self.software_manager.arp.get_arp_cache_mac_address(route.next_hop_ip_address)
outbound_network_interface = self.software_manager.arp.get_arp_cache_network_interface(
route.next_hop_ip_address
)
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
class Router(NetworkNode):
"""
Represents a network router, managing routing and forwarding of IP packets across network interfaces.
@@ -1049,6 +1214,10 @@ class Router(NetworkNode):
if not kwargs.get("route_table"):
kwargs["route_table"] = RouteTable(sys_log=kwargs["sys_log"])
super().__init__(hostname=hostname, num_ports=num_ports, **kwargs)
self.session_manager = RouterSessionManager(sys_log=self.sys_log)
self.session_manager.node = self
self.software_manager.session_manager = self.session_manager
self.session_manager.software_manager = self.software_manager
for i in range(1, self.num_ports + 1):
network_interface = RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0")
self.connect_nic(network_interface)
@@ -1068,8 +1237,7 @@ class Router(NetworkNode):
icmp: RouterICMP = self.software_manager.icmp # noqa
icmp.router = self
self.software_manager.install(RouterARP)
arp: RouterARP = self.software_manager.arp # noqa
arp.router = self
self.arp.router = self
def _set_default_acl(self):
"""
@@ -1313,6 +1481,8 @@ class Router(NetworkNode):
frame.ethernet.src_mac_addr = network_interface.mac_address
frame.ethernet.dst_mac_addr = target_mac
network_interface.send_frame(frame)
else:
self.sys_log.error(f"Frame dropped as there is no route to {frame.ip.dst_ip_address}")
def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]):
"""
@@ -1393,6 +1563,13 @@ class Router(NetworkNode):
- protocol (str, optional): the named IP protocol such as ICMP, TCP, or UDP
- src_ip_address (str, optional): IP address octet written in base 10
- dst_ip_address (str, optional): IP address octet written in base 10
- routes (list[dict]): List of route dicts with values:
- address (str): The destination address of the route.
- subnet_mask (str): The subnet mask of the route.
- next_hop_ip_address (str): The next hop IP for the route.
- metric (int): The metric of the route. Optional.
- default_route:
- next_hop_ip_address (str): The next hop IP for the route.
Example config:
```
@@ -1403,6 +1580,10 @@ class Router(NetworkNode):
1: {
'ip_address' : '192.168.1.1',
'subnet_mask' : '255.255.255.0',
},
2: {
'ip_address' : '192.168.0.1',
'subnet_mask' : '255.255.255.252',
}
},
'acl' : {
@@ -1410,6 +1591,10 @@ class Router(NetworkNode):
22: {'action': 'PERMIT', 'src_port': 'ARP', 'dst_port': 'ARP'},
23: {'action': 'PERMIT', 'protocol': 'ICMP'},
},
'routes' : [
{'address': '192.168.0.0', 'subnet_mask': '255.255.255.0', 'next_hop_ip_address': '192.168.1.2'}
],
'default_route': {'next_hop_ip_address': '192.168.0.2'}
}
```
@@ -1440,7 +1625,9 @@ class Router(NetworkNode):
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
if "routes" in cfg:
@@ -1451,4 +1638,8 @@ class Router(NetworkNode):
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
metric=float(route.get("metric", 0)),
)
if "default_route" in cfg:
next_hop_ip_address = cfg["default_route"].get("next_hop_ip_address", None)
if next_hop_ip_address:
router.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address)
return router

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

@@ -100,13 +100,8 @@ class Switch(NetworkNode):
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
for i in range(1, self.num_ports + 1):
self.connect_nic(SwitchPort())
def show(self, markdown: bool = False):
"""

71
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -1,10 +1,14 @@
from ipaddress import IPv4Address
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.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router, RouterInterface
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.utils.validators import IPV4Address
@@ -209,3 +213,68 @@ class WirelessRouter(Router):
raise NotImplementedError(
"Please use the 'configure_wireless_access_point' and 'configure_router_interface' functions."
)
@classmethod
def from_config(cls, cfg: Dict) -> "WirelessRouter":
"""Generate the wireless router from config.
Schema:
- hostname (str): unique name for this router.
- router_interface (dict): The values should be another dict specifying
- ip_address (str)
- subnet_mask (str)
- wireless_access_point (dict): Dict with
- ip address,
- subnet mask,
- frequency, (string: either WIFI_2_4 or WIFI_5)
- acl (dict): Dict with integers from 1 - max_acl_rules as keys. The key defines the position within the ACL
where the rule will be added (lower number is resolved first). The values should describe valid ACL
Rules as:
- action (str): either PERMIT or DENY
- src_port (str, optional): the named port such as HTTP, HTTPS, or POSTGRES_SERVER
- dst_port (str, optional): the named port such as HTTP, HTTPS, or POSTGRES_SERVER
- protocol (str, optional): the named IP protocol such as ICMP, TCP, or UDP
- src_ip_address (str, optional): IP address octet written in base 10
- dst_ip_address (str, optional): IP address octet written in base 10
:param cfg: Config dictionary
:type cfg: Dict
:return: WirelessRouter instance.
:rtype: WirelessRouter
"""
operating_state = (
NodeOperatingState.ON if not (p := cfg.get("operating_state")) else NodeOperatingState[p.upper()]
)
router = cls(hostname=cfg["hostname"], operating_state=operating_state)
if "router_interface" in cfg:
ip_address = cfg["router_interface"]["ip_address"]
subnet_mask = cfg["router_interface"]["subnet_mask"]
router.configure_router_interface(ip_address=ip_address, subnet_mask=subnet_mask)
if "wireless_access_point" in cfg:
ip_address = cfg["wireless_access_point"]["ip_address"]
subnet_mask = cfg["wireless_access_point"]["subnet_mask"]
frequency = AirSpaceFrequency[cfg["wireless_access_point"]["frequency"]]
router.configure_wireless_access_point(ip_address=ip_address, subnet_mask=subnet_mask, frequency=frequency)
if "acl" in cfg:
for r_num, r_cfg in cfg["acl"].items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
if "routes" in cfg:
for route in cfg.get("routes"):
router.route_table.add_route(
address=IPv4Address(route.get("address")),
subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")),
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
metric=float(route.get("metric", 0)),
)
return router

2
src/primaite/simulator/network/nmne.py
View File

@@ -6,7 +6,7 @@ CAPTURE_NMNE: bool = True
NMNE_CAPTURE_KEYWORDS: List[str] = []
"""List of keywords to identify malicious network events."""
# TODO: Remove final and make configurable after example layout when the NicObservation creates nmne structure dynamically
# TODO: Remove final and make configurable after example layout when the NICObservation creates nmne structure dynamically
CAPTURE_BY_DIRECTION: Final[bool] = True
"""Flag to determine if captures should be organized by traffic direction (inbound/outbound)."""
CAPTURE_BY_IP_ADDRESS: Final[bool] = False

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

@@ -8,7 +8,7 @@ from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.protocols.packet import DataPacket
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.transmission.transport_layer import Port, TCPHeader, UDPHeader
from primaite.simulator.network.utils import convert_bytes_to_megabits
_LOGGER = getLogger(__name__)
@@ -141,3 +141,37 @@ class Frame(BaseModel):
def size_Mbits(self) -> float: # noqa - Keep it as MBits as this is how they're expressed
"""The daa transfer size of the Frame in Mbits."""
return convert_bytes_to_megabits(self.size)
@property
def is_broadcast(self) -> bool:
"""
Determines if the Frame is a broadcast frame.
A Frame is considered a broadcast frame if the destination MAC address is set to the broadcast address
"ff:ff:ff:ff:ff:ff".
:return: True if the destination MAC address is a broadcast address, otherwise False.
"""
return self.ethernet.dst_mac_addr.lower() == "ff:ff:ff:ff:ff:ff"
@property
def is_arp(self) -> bool:
"""
Checks if the Frame is an ARP (Address Resolution Protocol) packet.
This is determined by checking if the destination port of the TCP header is equal to the ARP port.
:return: True if the Frame is an ARP packet, otherwise False.
"""
return self.udp.dst_port == Port.ARP
@property
def is_icmp(self) -> bool:
"""
Determines if the Frame is an ICMP (Internet Control Message Protocol) packet.
This check is performed by verifying if the 'icmp' attribute of the Frame instance is present (not None).
:return: True if the Frame is an ICMP packet (i.e., has an ICMP header), otherwise False.
"""
return self.icmp is not None

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

@@ -11,6 +11,9 @@ class Port(Enum):
.. _List of Ports:
"""
UNUSED = -1
"An unused port stub."
NONE = 0
"Place holder for a non-port."
WOL = 9

5
src/primaite/simulator/sim_container.py
View File

@@ -63,3 +63,8 @@ class Simulation(SimComponent):
"""Apply a timestep to the simulation."""
super().apply_timestep(timestep)
self.network.apply_timestep(timestep)
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
self.network.pre_timestep(timestep)

23
src/primaite/simulator/system/applications/application.py
View File

@@ -3,6 +3,8 @@ from enum import Enum
from typing import Any, Dict, Set
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.software import IOSoftware, SoftwareHealthState
_LOGGER = getLogger(__name__)
@@ -38,6 +40,17 @@ class Application(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def _init_request_manager(self) -> RequestManager:
"""
Initialise the request manager.
More information in user guide and docstring for SimComponent._init_request_manager.
"""
rm = super()._init_request_manager()
rm.add_request("close", RequestType(func=lambda request, context: RequestResponse.from_bool(self.close())))
return rm
@abstractmethod
def describe_state(self) -> Dict:
"""
@@ -67,7 +80,10 @@ class Application(IOSoftware):
"""
super().apply_timestep(timestep=timestep)
self.num_executions = 0 # reset number of executions
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
self.num_executions = 0
def _can_perform_action(self) -> bool:
"""
@@ -83,7 +99,7 @@ class Application(IOSoftware):
if self.operating_state is not self.operating_state.RUNNING:
# service is not running
_LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
_LOGGER.debug(f"Cannot perform action: {self.name} is {self.operating_state.name}")
return False
return True
@@ -104,11 +120,12 @@ class Application(IOSoftware):
"""The main application loop."""
pass
def close(self) -> None:
def close(self) -> bool:
"""Close the Application."""
if self.operating_state == ApplicationOperatingState.RUNNING:
self.sys_log.info(f"Closed Application{self.name}")
self.operating_state = ApplicationOperatingState.CLOSED
return True
def install(self) -> None:
"""Install Application."""

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

@@ -29,6 +29,9 @@ class DatabaseClient(Application):
_query_success_tracker: Dict[str, bool] = {}
_last_connection_successful: Optional[bool] = None
"""Keep track of connections that were established or verified during this step. Used for rewards."""
last_query_response: Optional[Dict] = None
"""Keep track of the latest query response. Used to determine rewards."""
_server_connection_id: Optional[str] = None
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
@@ -49,10 +52,9 @@ class DatabaseClient(Application):
def execute(self) -> bool:
"""Execution definition for db client: perform a select query."""
self.num_executions += 1 # trying to connect counts as an execution
if self.connections:
can_connect = self.check_connection(connection_id=list(self.connections.keys())[-1])
else:
can_connect = self.check_connection(connection_id=str(uuid4()))
if not self._server_connection_id:
self.connect()
can_connect = self.check_connection(connection_id=self._server_connection_id)
self._last_connection_successful = can_connect
return can_connect
@@ -78,17 +80,21 @@ class DatabaseClient(Application):
self.server_password = server_password
self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.")
def connect(self, connection_id: Optional[str] = None) -> bool:
def connect(self) -> bool:
"""Connect to a Database Service."""
if not self._can_perform_action():
return False
if not connection_id:
connection_id = str(uuid4())
if not self._server_connection_id:
self._server_connection_id = str(uuid4())
self.connected = self._connect(
server_ip_address=self.server_ip_address, password=self.server_password, connection_id=connection_id
server_ip_address=self.server_ip_address,
password=self.server_password,
connection_id=self._server_connection_id,
)
if not self.connected:
self._server_connection_id = None
return self.connected
def check_connection(self, connection_id: str) -> bool:
@@ -123,7 +129,7 @@ class DatabaseClient(Application):
:type: is_reattempt: Optional[bool]
"""
if is_reattempt:
if self.connections.get(connection_id):
if self._server_connection_id:
self.sys_log.info(
f"{self.name} {connection_id=}: DatabaseClient connection to {server_ip_address} authorised"
)
@@ -147,31 +153,28 @@ class DatabaseClient(Application):
server_ip_address=server_ip_address, password=password, connection_id=connection_id, is_reattempt=True
)
def disconnect(self, connection_id: Optional[str] = None) -> bool:
def disconnect(self) -> bool:
"""Disconnect from the Database Service."""
if not self._can_perform_action():
self.sys_log.error(f"Unable to disconnect - {self.name} is {self.operating_state.name}")
return False
# if there are no connections - nothing to disconnect
if not len(self.connections):
if not self._server_connection_id:
self.sys_log.error(f"Unable to disconnect - {self.name} has no active connections.")
return False
# if no connection provided, disconnect the first connection
if not connection_id:
connection_id = list(self.connections.keys())[0]
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "disconnect", "connection_id": connection_id},
payload={"type": "disconnect", "connection_id": self._server_connection_id},
dest_ip_address=self.server_ip_address,
dest_port=self.port,
)
self.remove_connection(connection_id=connection_id)
self.remove_connection(connection_id=self._server_connection_id)
self.sys_log.info(
f"{self.name}: DatabaseClient disconnected connection {connection_id} from {self.server_ip_address}"
f"{self.name}: DatabaseClient disconnected {self._server_connection_id} from {self.server_ip_address}"
)
self.connected = False
@@ -219,18 +222,23 @@ class DatabaseClient(Application):
if not self._can_perform_action():
return False
if connection_id is None:
if self.connections:
connection_id = list(self.connections.keys())[-1]
# TODO: if the most recent connection dies, it should be automatically cleared.
else:
connection_id = str(uuid4())
# reset last query response
self.last_query_response = None
if not self.connections.get(connection_id):
if not self.connect(connection_id=connection_id):
return False
connection_id: str
if not connection_id:
connection_id = self._server_connection_id
if not connection_id:
self.connect()
connection_id = self._server_connection_id
if not connection_id:
msg = "Cannot run sql query, could not establish connection with the server."
self.parent.sys_log.error(msg)
return False
# Initialise the tracker of this ID to False
uuid = str(uuid4())
self._query_success_tracker[uuid] = False
return self._query(sql=sql, query_id=uuid, connection_id=connection_id)
@@ -252,6 +260,7 @@ class DatabaseClient(Application):
# add connection
self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id)
elif payload["type"] == "sql":
self.last_query_response = payload
query_id = payload.get("uuid")
status_code = payload.get("status_code")
self._query_success_tracker[query_id] = status_code == 200

316
src/primaite/simulator/system/applications/red_applications/ransomware_script.py Normal file
View File

@@ -0,0 +1,316 @@
from enum import IntEnum
from ipaddress import IPv4Address
from typing import Dict, Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
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
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
class RansomwareAttackStage(IntEnum):
"""
Enumeration representing different attack stages of the ransomware script.
This enumeration defines the various stages a data manipulation attack can be in during its lifecycle
in the simulation.
Each stage represents a specific phase in the attack process.
"""
NOT_STARTED = 0
"Indicates that the attack has not started yet."
DOWNLOAD = 1
"Installing the Encryption Script - Testing"
INSTALL = 2
"The stage where logon procedures are simulated."
ACTIVATE = 3
"Operating Status Changes"
PROPAGATE = 4
"Represents the stage of performing a horizontal port scan on the target."
COMMAND_AND_CONTROL = 5
"Represents the stage of setting up a rely C2 Beacon (Not Implemented)"
PAYLOAD = 6
"Stage of actively attacking the target."
SUCCEEDED = 7
"Indicates the attack has been successfully completed."
FAILED = 8
"Signifies that the attack has failed."
class RansomwareScript(Application):
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
:ivar payload: The attack stage query payload. (Default Corrupt)
:ivar target_scan_p_of_success: The probability of success for the target scan stage.
:ivar c2_beacon_p_of_success: The probability of success for the c2_beacon stage
:ivar ransomware_encrypt_p_of_success: The probability of success for the ransomware 'attack' (encrypt) stage.
:ivar repeat: Whether to repeat attacking once finished.
"""
server_ip_address: Optional[IPv4Address] = None
"""IP address of node which hosts the database."""
server_password: Optional[str] = None
"""Password required to access the database."""
payload: Optional[str] = "ENCRYPT"
"Payload String for the payload stage"
target_scan_p_of_success: float = 0.9
"Probability of the target scan succeeding: Default 0.9"
c2_beacon_p_of_success: float = 0.9
"Probability of the c2 beacon setup stage succeeding: Default 0.9"
ransomware_encrypt_p_of_success: float = 0.9
"Probability of the ransomware attack succeeding: Default 0.9"
repeat: bool = False
"If true, the Denial of Service bot will keep performing the attack."
attack_stage: RansomwareAttackStage = RansomwareAttackStage.NOT_STARTED
"The ransomware attack stage. See RansomwareAttackStage Class"
def __init__(self, **kwargs):
kwargs["name"] = "RansomwareScript"
kwargs["port"] = Port.NONE
kwargs["protocol"] = IPProtocol.NONE
super().__init__(**kwargs)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: Current state of this object and child objects.
:rtype: Dict
"""
state = super().describe_state()
return state
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the Ransomware Script."""
db_client = self.software_manager.software.get("DatabaseClient")
if db_client is None:
_LOGGER.info(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client
def _init_request_manager(self) -> RequestManager:
"""
Initialise the request manager.
More information in user guide and docstring for SimComponent._init_request_manager.
"""
rm = super()._init_request_manager()
rm.add_request(
name="execute",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.attack())),
)
return rm
def _activate(self):
"""
Simulate the install process as the initial stage of the attack.
Advances the attack stage to 'ACTIVATE' attack state.
"""
if self.attack_stage == RansomwareAttackStage.INSTALL:
self.sys_log.info(f"{self.name}: Activated!")
self.attack_stage = RansomwareAttackStage.ACTIVATE
def apply_timestep(self, timestep: int) -> None:
"""
Apply a timestep to the bot, triggering the application loop.
:param timestep: The timestep value to update the bot's state.
"""
pass
def run(self) -> bool:
"""Calls the parent classes execute method before starting the application loop."""
super().run()
return True
def _application_loop(self) -> bool:
"""
The main application loop of the script, handling the attack process.
This is the core loop where the bot sequentially goes through the stages of the attack.
"""
if not self._can_perform_action():
return False
if self.server_ip_address and self.payload:
self.sys_log.info(f"{self.name}: Running")
self.attack_stage = RansomwareAttackStage.NOT_STARTED
self._local_download()
self._install()
self._activate()
self._perform_target_scan()
self._setup_beacon()
self._perform_ransomware_encrypt()
if self.repeat and self.attack_stage in (
RansomwareAttackStage.SUCCEEDED,
RansomwareAttackStage.FAILED,
):
self.attack_stage = RansomwareAttackStage.NOT_STARTED
return True
else:
self.sys_log.error(f"{self.name}: Failed to start as it requires both a target_ip_address and payload.")
return False
def configure(
self,
server_ip_address: IPv4Address,
server_password: Optional[str] = None,
payload: Optional[str] = None,
target_scan_p_of_success: Optional[float] = None,
c2_beacon_p_of_success: Optional[float] = None,
ransomware_encrypt_p_of_success: Optional[float] = None,
repeat: bool = True,
):
"""
Configure the Ransomware Script to communicate with a DatabaseService.
:param server_ip_address: The IP address of the Node the DatabaseService is on.
:param server_password: The password on the DatabaseService.
:param payload: The attack stage query (Encrypt / Delete)
:param target_scan_p_of_success: The probability of success for the target scan stage.
:param c2_beacon_p_of_success: The probability of success for the c2_beacon stage
:param ransomware_encrypt_p_of_success: The probability of success for the ransomware 'attack' (encrypt) stage.
:param repeat: Whether to repeat attacking once finished.
"""
if server_ip_address:
self.server_ip_address = server_ip_address
if server_password:
self.server_password = server_password
if payload:
self.payload = payload
if target_scan_p_of_success:
self.target_scan_p_of_success = target_scan_p_of_success
if c2_beacon_p_of_success:
self.c2_beacon_p_of_success = c2_beacon_p_of_success
if ransomware_encrypt_p_of_success:
self.ransomware_encrypt_p_of_success = ransomware_encrypt_p_of_success
if repeat:
self.repeat = repeat
self.sys_log.info(
f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
f"{repeat=}."
)
def _install(self):
"""
Simulate the install stage in the kill-chain.
Advances the attack stage to 'ACTIVATE' if successful.
From this attack stage onwards.
the ransomware application is now visible from this point onwardin the observation space.
"""
if self.attack_stage == RansomwareAttackStage.DOWNLOAD:
self.sys_log.info(f"{self.name}: Malware installed on the local file system")
downloads_folder = self.file_system.get_folder(folder_name="downloads")
ransomware_file = downloads_folder.get_file(file_name="ransom_script.pdf")
ransomware_file.num_access += 1
self.attack_stage = RansomwareAttackStage.INSTALL
def _setup_beacon(self):
"""
Simulates setting up a c2 beacon; currently a pseudo step for increasing red variance.
Advances the attack stage to 'COMMAND AND CONTROL` if successful.
:param p_of_sucess: Probability of a successful c2 setup (Advancing this step),
by default the success rate is 0.5
"""
if self.attack_stage == RansomwareAttackStage.PROPAGATE:
self.sys_log.info(f"{self.name} Attempting to set up C&C Beacon - Scan 1/2")
if simulate_trial(self.c2_beacon_p_of_success):
self.sys_log.info(f"{self.name} C&C Successful setup - Scan 2/2")
c2c_setup = True # TODO Implement the c2c step via an FTP Application/Service
if c2c_setup:
self.attack_stage = RansomwareAttackStage.COMMAND_AND_CONTROL
def _perform_target_scan(self):
"""
Perform a simulated port scan to check for open SQL ports.
Advances the attack stage to `PROPAGATE` if successful.
:param p_of_success: Probability of successful port scan, by default 0.1.
"""
if self.attack_stage == RansomwareAttackStage.ACTIVATE:
# perform a port scan to identify that the SQL port is open on the server
self.sys_log.info(f"{self.name}: Scanning for vulnerable databases - Scan 0/2")
if simulate_trial(self.target_scan_p_of_success):
self.sys_log.info(f"{self.name}: Found a target database! Scan 1/2")
port_is_open = True # TODO Implement a NNME Triggering scan as a seperate Red Application
if port_is_open:
self.attack_stage = RansomwareAttackStage.PROPAGATE
def attack(self) -> bool:
"""Perform the attack steps after opening the application."""
if not self._can_perform_action():
_LOGGER.debug("Ransomware application is unable to perform it's actions.")
self.run()
self.num_executions += 1
return self._application_loop()
def _perform_ransomware_encrypt(self):
"""
Execute the Ransomware Encrypt payload on the target.
Advances the attack stage to `COMPLETE` if successful, or 'FAILED' if unsuccessful.
:param p_of_success: Probability of successfully performing ransomware encryption, by default 0.1.
"""
if self._host_db_client is None:
self.sys_log.info(f"{self.name}: Failed to connect to db_client - Ransomware Script")
self.attack_stage = RansomwareAttackStage.FAILED
return
self._host_db_client.server_ip_address = self.server_ip_address
self._host_db_client.server_password = self.server_password
if self.attack_stage == RansomwareAttackStage.COMMAND_AND_CONTROL:
if simulate_trial(self.ransomware_encrypt_p_of_success):
self.sys_log.info(f"{self.name}: Attempting to launch payload")
if not len(self._host_db_client.connections):
self._host_db_client.connect()
if len(self._host_db_client.connections):
self._host_db_client.query(self.payload)
self.sys_log.info(f"{self.name} Payload delivered: {self.payload}")
attack_successful = True
if attack_successful:
self.sys_log.info(f"{self.name}: Payload Successful")
self.attack_stage = RansomwareAttackStage.SUCCEEDED
else:
self.sys_log.info(f"{self.name}: Payload failed")
self.attack_stage = RansomwareAttackStage.FAILED
else:
self.sys_log.error("Attack Attempted to launch too quickly")
self.attack_stage = RansomwareAttackStage.FAILED
def _local_download(self):
"""Downloads itself via the onto the local file_system."""
if self.attack_stage == RansomwareAttackStage.NOT_STARTED:
if self._local_download_verify():
self.attack_stage = RansomwareAttackStage.DOWNLOAD
else:
self.sys_log.info("Malware failed to create a installation location")
self.attack_stage = RansomwareAttackStage.FAILED
else:
self.sys_log.info("Malware failed to download")
self.attack_stage = RansomwareAttackStage.FAILED
def _local_download_verify(self) -> bool:
"""Verifies a download location - Creates one if needed."""
for folder in self.file_system.folders:
if self.file_system.folders[folder].name == "downloads":
self.file_system.num_file_creations += 1
return True
self.file_system.create_folder("downloads")
self.file_system.create_file(folder_name="downloads", file_name="ransom_script.pdf")
return True

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

@@ -49,8 +49,9 @@ class PacketCapture:
self.current_episode: int = 1
self.setup_logger(outbound=False)
self.setup_logger(outbound=True)
if SIM_OUTPUT.save_pcap_logs:
self.setup_logger(outbound=False)
self.setup_logger(outbound=True)
def setup_logger(self, outbound: bool = False):
"""Set up the logger configuration."""
@@ -108,8 +109,9 @@ class PacketCapture:
:param frame: The PCAP frame to capture.
"""
msg = frame.model_dump_json()
self.inbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL
if SIM_OUTPUT.save_pcap_logs:
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 ;(
"""
@@ -117,5 +119,6 @@ class PacketCapture:
: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
if SIM_OUTPUT.save_pcap_logs:
msg = frame.model_dump_json()
self.outbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL

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

@@ -72,7 +72,6 @@ class SessionManager:
Manages network sessions, including session creation, lookup, and communication with other components.
:param sys_log: A reference to the system log component.
:param arp_cache: A reference to the ARP cache component.
"""
def __init__(self, sys_log: SysLog):

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

@@ -88,6 +88,10 @@ class SysLog:
root.mkdir(exist_ok=True, parents=True)
return root / f"{self.hostname}_sys.log"
def _write_to_terminal(self, msg: str, level: str, to_terminal: bool = False):
if to_terminal or SIM_OUTPUT.write_sys_log_to_terminal:
print(f"{self.hostname}: ({level}) {msg}")
def debug(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the DEBUG level.
@@ -97,8 +101,7 @@ class SysLog:
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.debug(msg)
if to_terminal:
print(msg)
self._write_to_terminal(msg, "DEBUG", to_terminal)
def info(self, msg: str, to_terminal: bool = False):
"""
@@ -109,8 +112,7 @@ class SysLog:
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.info(msg)
if to_terminal:
print(msg)
self._write_to_terminal(msg, "INFO", to_terminal)
def warning(self, msg: str, to_terminal: bool = False):
"""
@@ -121,8 +123,7 @@ class SysLog:
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.warning(msg)
if to_terminal:
print(msg)
self._write_to_terminal(msg, "WARNING", to_terminal)
def error(self, msg: str, to_terminal: bool = False):
"""
@@ -133,8 +134,7 @@ class SysLog:
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.error(msg)
if to_terminal:
print(msg)
self._write_to_terminal(msg, "ERROR", to_terminal)
def critical(self, msg: str, to_terminal: bool = False):
"""
@@ -145,5 +145,4 @@ class SysLog:
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.critical(msg)
if to_terminal:
print(msg)
self._write_to_terminal(msg, "CRITICAL", to_terminal)

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

@@ -65,6 +65,10 @@ class ARP(Service):
"""Clears the arp cache."""
self.arp.clear()
def get_default_gateway_network_interface(self) -> Optional[NetworkInterface]:
"""Not used at the parent ARP level. Should return None when there is no override by child class."""
return None
def add_arp_cache_entry(
self, ip_address: IPV4Address, mac_address: str, network_interface: NetworkInterface, override: bool = False
):

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

@@ -104,14 +104,30 @@ class DatabaseService(Service):
self.sys_log.error("Unable to restore database backup.")
return False
old_visible_state = SoftwareHealthState.GOOD
# get db file regardless of whether or not it was deleted
db_file = self.file_system.get_file(folder_name="database", file_name="database.db", include_deleted=True)
if db_file is None:
self.sys_log.error("Database file not initialised.")
return False
# if the file was deleted, get the old visible health state
if db_file.deleted:
old_visible_state = db_file.visible_health_status
else:
old_visible_state = self.db_file.visible_health_status
self.file_system.delete_file(folder_name="database", file_name="database.db")
# replace db file
self.file_system.delete_file(folder_name="database", file_name="database.db")
self.file_system.copy_file(src_folder_name="downloads", src_file_name="database.db", dst_folder_name="database")
if self.db_file is None:
self.sys_log.error("Copying database backup failed.")
return False
self.db_file.visible_health_status = old_visible_state
self.set_health_state(SoftwareHealthState.GOOD)
return True
@@ -125,8 +141,7 @@ class DatabaseService(Service):
"""Returns the database file."""
return self.file_system.get_file(folder_name="database", file_name="database.db")
@property
def folder(self) -> Folder:
def _return_database_folder(self) -> Folder:
"""Returns the database folder."""
return self.file_system.get_folder_by_id(self.db_file.folder_id)
@@ -171,7 +186,10 @@ class DatabaseService(Service):
}
def _process_sql(
self, query: Literal["SELECT", "DELETE", "INSERT"], query_id: str, connection_id: Optional[str] = None
self,
query: Literal["SELECT", "DELETE", "INSERT", "ENCRYPT"],
query_id: str,
connection_id: Optional[str] = None,
) -> Dict[str, Union[int, List[Any]]]:
"""
Executes the given SQL query and returns the result.
@@ -180,6 +198,7 @@ class DatabaseService(Service):
- SELECT : returns the data
- DELETE : deletes the data
- INSERT : inserts the data
- ENCRYPT : corrupts the data
:param query: The SQL query to be executed.
:return: Dictionary containing status code and data fetched.
@@ -188,10 +207,18 @@ class DatabaseService(Service):
if not self.db_file:
self.sys_log.info(f"{self.name}: Failed to run {query} because the database file is missing.")
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
if query == "SELECT":
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
if self.db_file.health_status == FileSystemItemHealthStatus.CORRUPT:
return {
"status_code": 200,
"type": "sql",
"data": False,
"uuid": query_id,
"connection_id": connection_id,
}
elif self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {
"status_code": 200,
"type": "sql",
@@ -200,7 +227,7 @@ class DatabaseService(Service):
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
elif query == "DELETE":
self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
return {
@@ -210,6 +237,20 @@ class DatabaseService(Service):
"uuid": query_id,
"connection_id": connection_id,
}
elif query == "ENCRYPT":
self.file_system.num_file_creations += 1
self.db_file.health_status = FileSystemItemHealthStatus.CORRUPT
self.db_file.num_access += 1
database_folder = self._return_database_folder()
database_folder.health_status = FileSystemItemHealthStatus.CORRUPT
self.file_system.num_file_deletions += 1
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 {
@@ -220,7 +261,7 @@ class DatabaseService(Service):
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
elif query == "SELECT * FROM pg_stat_activity":
# Check if the connection is active.
if self.health_state_actual == SoftwareHealthState.GOOD:
@@ -304,8 +345,8 @@ class DatabaseService(Service):
self.backup_database()
return super().apply_timestep(timestep)
def _update_patch_status(self) -> None:
"""Perform a database restore when the patching countdown is finished."""
super()._update_patch_status()
if self._patching_countdown is None:
def _update_fix_status(self) -> None:
"""Perform a database restore when the FIXING countdown is finished."""
super()._update_fix_status()
if self._fixing_countdown is None:
self.restore_backup()

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

@@ -87,13 +87,9 @@ class NTPClient(Service):
:return: True if successful, False otherwise.
"""
if not isinstance(payload, NTPPacket):
_LOGGER.debug(f"{payload} is not a NTPPacket")
_LOGGER.debug(f"{self.name}: Failed to parse NTP update")
return False
if payload.ntp_reply.ntp_datetime:
self.sys_log.info(
f"{self.name}: \
Received time update from NTP server{payload.ntp_reply.ntp_datetime}"
)
self.time = payload.ntp_reply.ntp_datetime
return True
@@ -124,5 +120,3 @@ class NTPClient(Service):
if self.operating_state == ServiceOperatingState.RUNNING:
# request time from server
self.request_time()
else:
self.sys_log.debug(f"{self.name} ntp client not running")

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

@@ -59,7 +59,7 @@ class Service(IOSoftware):
if self.operating_state is not ServiceOperatingState.RUNNING:
# service is not running
_LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
_LOGGER.debug(f"Cannot perform action: {self.name} is {self.operating_state.name}")
return False
return True

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

@@ -43,8 +43,8 @@ class SoftwareHealthState(Enum):
"Unused state."
GOOD = 1
"The software is in a good and healthy condition."
PATCHING = 2
"The software is undergoing patching or updates."
FIXING = 2
"The software is undergoing FIXING or updates."
COMPROMISED = 3
"The software's security has been compromised."
OVERWHELMED = 4
@@ -82,13 +82,13 @@ class Software(SimComponent):
"The health state of the software visible to the red agent."
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
"The criticality level of the software."
patching_count: int = 0
fixing_count: int = 0
"The count of patches applied to the software, defaults to 0."
scanning_count: int = 0
"The count of times the software has been scanned, defaults to 0."
revealed_to_red: bool = False
"Indicates if the software has been revealed to red agent, defaults is False."
software_manager: "SoftwareManager" = None
software_manager: Optional["SoftwareManager"] = None
"An instance of Software Manager that is used by the parent node."
sys_log: SysLog = None
"An instance of SysLog that is used by the parent node."
@@ -96,9 +96,9 @@ class Software(SimComponent):
"The FileSystem of the Node the Software is installed on."
folder: Optional[Folder] = None
"The folder on the file system the Software uses."
patching_duration: int = 2
fixing_duration: int = 2
"The number of ticks it takes to patch the software."
_patching_countdown: Optional[int] = None
_fixing_countdown: Optional[int] = None
"Current number of ticks left to patch the software."
def _init_request_manager(self) -> RequestManager:
@@ -117,9 +117,9 @@ class Software(SimComponent):
),
)
rm.add_request(
"patch",
"fix",
RequestType(
func=lambda request, context: RequestResponse.from_bool(self.patch()),
func=lambda request, context: RequestResponse.from_bool(self.fix()),
),
)
rm.add_request("scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan())))
@@ -149,7 +149,7 @@ class Software(SimComponent):
"health_state_actual": self.health_state_actual.value,
"health_state_visible": self.health_state_visible.value,
"criticality": self.criticality.value,
"patching_count": self.patching_count,
"fixing_count": self.fixing_count,
"scanning_count": self.scanning_count,
"revealed_to_red": self.revealed_to_red,
}
@@ -194,21 +194,21 @@ class Software(SimComponent):
self.health_state_visible = self.health_state_actual
return True
def patch(self) -> bool:
"""Perform a patch on the software."""
def fix(self) -> bool:
"""Perform a fix on the software."""
if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD):
self._patching_countdown = self.patching_duration
self.set_health_state(SoftwareHealthState.PATCHING)
self._fixing_countdown = self.fixing_duration
self.set_health_state(SoftwareHealthState.FIXING)
return True
return False
def _update_patch_status(self) -> None:
"""Update the patch status of the software."""
self._patching_countdown -= 1
if self._patching_countdown <= 0:
def _update_fix_status(self) -> None:
"""Update the fix status of the software."""
self._fixing_countdown -= 1
if self._fixing_countdown <= 0:
self.set_health_state(SoftwareHealthState.GOOD)
self._patching_countdown = None
self.patching_count += 1
self._fixing_countdown = None
self.fixing_count += 1
def reveal_to_red(self) -> None:
"""Reveals the software to the red agent."""
@@ -221,8 +221,12 @@ class Software(SimComponent):
:param timestep: The current timestep of the simulation.
"""
super().apply_timestep(timestep)
if self.health_state_actual == SoftwareHealthState.PATCHING:
self._update_patch_status()
if self.health_state_actual == SoftwareHealthState.FIXING:
self._update_fix_status()
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
class IOSoftware(Software):