From 4d3c85bc141464eef5ab592c7d7994705ac4b07b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 7 May 2024 15:35:02 +0100 Subject: [PATCH] #2550 Backport changes into core --- .../notebooks/Training-an-RLLib-Agent.ipynb | 14 +- .../create-simulation_demo.ipynb | 2 +- .../network_simulator_demo.ipynb | 218 ++++++++++-------- src/primaite/simulator/network/airspace.py | 9 +- .../simulator/network/hardware/base.py | 2 +- .../system/applications/application.py | 13 +- .../applications/red_applications/dos_bot.py | 1 + .../red_applications/ransomware_script.py | 8 - .../test_uc2_data_manipulation_scenario.py | 5 + 9 files changed, 151 insertions(+), 121 deletions(-) diff --git a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb index 2fe84655..60737ee5 100644 --- a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb @@ -45,15 +45,9 @@ "metadata": {}, "outputs": [], "source": [ - "print(cfg['agents'][2]['agent_settings'])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ + "for agent in cfg['agents']:\n", + " if agent[\"ref\"] == \"defender\":\n", + " agent['agent_settings']['flatten_obs'] = True\n", "env_config = cfg\n", "\n", "config = (\n", @@ -80,7 +74,7 @@ "tune.Tuner(\n", " \"PPO\",\n", " run_config=air.RunConfig(\n", - " stop={\"timesteps_total\": 512}\n", + " stop={\"timesteps_total\": 1e3 * 128}\n", " ),\n", " param_space=config\n", ").fit()\n" diff --git a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb index 06ecd4be..31173022 100644 --- a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb +++ b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb @@ -261,7 +261,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb index 7f4cf3b1..107a2565 100644 --- a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb +++ b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb @@ -256,9 +256,11 @@ { "cell_type": "markdown", "id": "22", - "metadata": {}, + "metadata": { + "tags": [] + }, "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=`." + "Calling `switch.arp.show()` displays the Switch ARP Cache." ] }, { @@ -270,13 +272,33 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"switch_1\").sys_log.show()" + "network.get_node_by_hostname(\"switch_1\").arp.show()" ] }, { "cell_type": "markdown", "id": "24", "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=`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"switch_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, "source": [ "### Computer/Server Nodes\n", "\n", @@ -285,7 +307,7 @@ }, { "cell_type": "markdown", - "id": "25", + "id": "27", "metadata": { "tags": [] }, @@ -293,26 +315,6 @@ "Calling `computer.show()` displays the NICs on the Computer/Server." ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "26", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"security_suite\").show()" - ] - }, - { - "cell_type": "markdown", - "id": "27", - "metadata": {}, - "source": [ - "Calling `computer.arp.show()` displays the Computer/Server ARP Cache." - ] - }, { "cell_type": "code", "execution_count": null, @@ -322,7 +324,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"security_suite\").arp.show()" + "network.get_node_by_hostname(\"security_suite\").show()" ] }, { @@ -330,7 +332,7 @@ "id": "29", "metadata": {}, "source": [ - "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=`." + "Calling `computer.arp.show()` displays the Computer/Server ARP Cache." ] }, { @@ -342,7 +344,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"security_suite\").sys_log.show()" + "network.get_node_by_hostname(\"security_suite\").arp.show()" ] }, { @@ -350,9 +352,7 @@ "id": "31", "metadata": {}, "source": [ - "## Basic Network Comms Check\n", - "\n", - "We can perform a good old ping to check that Nodes are able to communicate with each other." + "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=`." ] }, { @@ -364,7 +364,7 @@ }, "outputs": [], "source": [ - "network.show(nodes=False, links=False)" + "network.get_node_by_hostname(\"security_suite\").sys_log.show()" ] }, { @@ -372,7 +372,9 @@ "id": "33", "metadata": {}, "source": [ - "We'll first ping client_1's default gateway." + "## Basic Network Comms Check\n", + "\n", + "We can perform a good old ping to check that Nodes are able to communicate with each other." ] }, { @@ -384,27 +386,27 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")" + "network.show(nodes=False, links=False)" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "We'll first ping client_1's default gateway." ] }, { "cell_type": "code", "execution_count": null, - "id": "35", + "id": "36", "metadata": { "tags": [] }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").sys_log.show(15)" - ] - }, - { - "cell_type": "markdown", - "id": "36", - "metadata": {}, - "source": [ - "Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)." + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")" ] }, { @@ -416,7 +418,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")" + "network.get_node_by_hostname(\"client_1\").sys_log.show(15)" ] }, { @@ -424,7 +426,7 @@ "id": "38", "metadata": {}, "source": [ - "And finally, we'll ping the web server." + "Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)." ] }, { @@ -436,7 +438,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")" ] }, { @@ -444,7 +446,7 @@ "id": "40", "metadata": {}, "source": [ - "To confirm that the ping was received and processed by the web_server, we can view the sys log" + "And finally, we'll ping the web server." ] }, { @@ -456,45 +458,45 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"web_server\").sys_log.show()" + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" ] }, { "cell_type": "markdown", "id": "42", "metadata": {}, + "source": [ + "To confirm that the ping was received and processed by the web_server, we can view the sys log" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"web_server\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "44", + "metadata": {}, "source": [ "## Advanced Network Usage\n", "\n", "We can now use the Network to perform some more advanced things." ] }, - { - "cell_type": "markdown", - "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..." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "44", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" - ] - }, { "cell_type": "markdown", "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:" + "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..." ] }, { @@ -506,13 +508,33 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_2\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" ] }, { "cell_type": "markdown", "id": "47", "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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "49", + "metadata": {}, "source": [ "Now we'll add an ACL to block ICMP from 192.168.10.22" ] @@ -520,7 +542,7 @@ { "cell_type": "code", "execution_count": null, - "id": "48", + "id": "50", "metadata": { "tags": [] }, @@ -540,7 +562,7 @@ { "cell_type": "code", "execution_count": null, - "id": "49", + "id": "51", "metadata": { "tags": [] }, @@ -549,32 +571,12 @@ "network.get_node_by_hostname(\"router_1\").acl.show()" ] }, - { - "cell_type": "markdown", - "id": "50", - "metadata": {}, - "source": [ - "Now we attempt (and fail) to ping the web server" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "51", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" - ] - }, { "cell_type": "markdown", "id": "52", "metadata": {}, "source": [ - "We can check that the ping was actually sent by client_2 by viewing the sys log" + "Now we attempt (and fail) to ping the web server" ] }, { @@ -586,7 +588,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_2\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" ] }, { @@ -594,7 +596,7 @@ "id": "54", "metadata": {}, "source": [ - "We can check the router sys log to see why the traffic was blocked" + "We can check that the ping was actually sent by client_2 by viewing the sys log" ] }, { @@ -606,7 +608,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"router_1\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").sys_log.show()" ] }, { @@ -614,7 +616,7 @@ "id": "56", "metadata": {}, "source": [ - "Now a final check to ensure that client_1 can still ping the web_server." + "We can check the router sys log to see why the traffic was blocked" ] }, { @@ -625,6 +627,26 @@ "tags": [] }, "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "58", + "metadata": {}, + "source": [ + "Now a final check to ensure that client_1 can still ping the web_server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59", + "metadata": { + "tags": [] + }, + "outputs": [], "source": [ "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" ] @@ -632,7 +654,7 @@ { "cell_type": "code", "execution_count": null, - "id": "58", + "id": "60", "metadata": { "tags": [] }, diff --git a/src/primaite/simulator/network/airspace.py b/src/primaite/simulator/network/airspace.py index 8a00a4a4..907ab233 100644 --- a/src/primaite/simulator/network/airspace.py +++ b/src/primaite/simulator/network/airspace.py @@ -192,7 +192,6 @@ class WirelessNetworkInterface(NetworkInterface, ABC): # Cannot send Frame as the network interface is not enabled return False - @abstractmethod def receive_frame(self, frame: Frame) -> bool: """ Receives a network frame on the network interface. @@ -200,7 +199,13 @@ class WirelessNetworkInterface(NetworkInterface, ABC): :param frame: The network frame being received. :return: A boolean indicating whether the frame was successfully received. """ - pass + if self.enabled: + frame.set_sent_timestamp() + self.pcap.capture_inbound(frame) + self._connected_node.receive_frame(frame, self) + return True + # Cannot receive Frame as the network interface is not enabled + return False class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC): diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 78a0bf2d..b4d32dc4 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1378,7 +1378,7 @@ class Node(SimComponent): application_instance.configure(server_ip_address=IPv4Address(ip_address)) else: pass - + application_instance.install() if application_instance.name in self.software_manager.software: return True else: diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 26603b43..294de27b 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -1,6 +1,6 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, Set +from typing import Any, Dict, Optional, Set from primaite.interface.request import RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -33,6 +33,10 @@ class Application(IOSoftware): "The number of times the application has been executed. Default is 0." groups: Set[str] = set() "The set of groups to which the application belongs." + install_duration: int = 2 + "How long it takes to install the application." + install_countdown: Optional[int] = None + "The countdown to the end of the installation process. None if not currently installing" def __init__(self, **kwargs): super().__init__(**kwargs) @@ -76,6 +80,12 @@ class Application(IOSoftware): :param timestep: The current timestep of the simulation. """ super().apply_timestep(timestep=timestep) + if self.operating_state is ApplicationOperatingState.INSTALLING: + self.install_countdown -= 1 + if self.install_countdown <= 0: + self.operating_state = ApplicationOperatingState.RUNNING + self.health_state_actual = SoftwareHealthState.GOOD + self.install_countdown = None def pre_timestep(self, timestep: int) -> None: """Apply pre-timestep logic.""" @@ -129,6 +139,7 @@ class Application(IOSoftware): super().install() if self.operating_state == ApplicationOperatingState.CLOSED: self.operating_state = ApplicationOperatingState.INSTALLING + self.install_countdown = self.install_duration def receive(self, payload: Any, session_id: str, **kwargs) -> bool: """ diff --git a/src/primaite/simulator/system/applications/red_applications/dos_bot.py b/src/primaite/simulator/system/applications/red_applications/dos_bot.py index 53fc9740..0e45aad9 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -177,4 +177,5 @@ class DoSBot(DatabaseClient): :param timestep: The timestep value to update the bot's state. """ + super().apply_timestep(timestep=timestep) self._application_loop() diff --git a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py index 4c2d7927..8acc07b4 100644 --- a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py +++ b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py @@ -118,14 +118,6 @@ class RansomwareScript(Application): 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() diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index 4e203669..5df8f964 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -62,6 +62,11 @@ def test_application_install_uninstall_on_uc2(): _, _, _, _, info = env.step(78) assert "DoSBot" in domcon.software_manager.software + # installing takes 3 steps so let's wait for 3 steps + env.step(0) + env.step(0) + env.step(0) + # Test we can now execute the DoSBot app _, _, _, _, info = env.step(81) assert info["agent_actions"]["defender"].response.status == "success"