From c34cb6d7ce6142a1989532d85240209e50ac06b8 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 1 Jul 2024 11:31:27 +0100 Subject: [PATCH 1/6] #2700 Add DatabaseConfigure action --- src/primaite/game/agent/actions.py | 23 +++++ .../system/applications/database_client.py | 18 +++- .../observations/actions/__init__.py | 1 + .../actions/test_configure_actions.py | 85 +++++++++++++++++++ 4 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 tests/integration_tests/game_layer/observations/actions/__init__.py create mode 100644 tests/integration_tests/game_layer/observations/actions/test_configure_actions.py diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index e165c9ad..4cb31d25 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -17,6 +17,7 @@ from gymnasium import spaces from pydantic import BaseModel, Field, field_validator, ValidationInfo from primaite import getLogger +from primaite.interface.request import RequestFormat _LOGGER = getLogger(__name__) @@ -245,6 +246,27 @@ class NodeApplicationInstallAction(AbstractAction): ] +class ConfigureDatabaseClientAction(AbstractAction): + """Action which sets config parameters for a database client on a node.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this action.""" + + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, options: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + ConfigureDatabaseClientAction._Opts.model_validate(options) # check that options adhere to schema + return ["network", "node", node_name, "application", "DatabaseClient", "configure", options] + + class NodeApplicationRemoveAction(AbstractAction): """Action which removes/uninstalls an application.""" @@ -1045,6 +1067,7 @@ class ActionManager: "NODE_NMAP_PING_SCAN": NodeNMAPPingScanAction, "NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction, "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction, + "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, } """Dictionary which maps action type strings to the corresponding action class.""" diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index bae2139b..6396c678 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -8,13 +8,14 @@ from uuid import uuid4 from prettytable import MARKDOWN, PrettyTable from pydantic import BaseModel -from primaite.interface.request import RequestResponse +from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.network.hardware.nodes.host.host_node import HostNode 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.core.software_manager import SoftwareManager +from primaite.utils.validators import IPV4Address class DatabaseClientConnection(BaseModel): @@ -96,6 +97,14 @@ class DatabaseClient(Application): """ rm = super()._init_request_manager() rm.add_request("execute", RequestType(func=lambda request, context: RequestResponse.from_bool(self.execute()))) + + def _configure(request: RequestFormat, context: Dict) -> RequestResponse: + ip, pw = request[-1].get("server_ip_address"), request[-1].get("server_password") + ip = None if ip is None else IPV4Address(ip) + success = self.configure(server_ip_address=ip, server_password=pw) + return RequestResponse.from_bool(success) + + rm.add_request("configure", RequestType(func=lambda request, context: _configure(request, context))) return rm def execute(self) -> bool: @@ -141,16 +150,17 @@ class DatabaseClient(Application): table.add_row([connection_id, connection.is_active]) print(table.get_string(sortby="Connection ID")) - def configure(self, server_ip_address: IPv4Address, server_password: Optional[str] = None): + def configure(self, server_ip_address: Optional[IPv4Address] = None, server_password: Optional[str] = None) -> bool: """ Configure the DatabaseClient 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. """ - self.server_ip_address = server_ip_address - self.server_password = server_password + self.server_ip_address = server_ip_address or self.server_ip_address + self.server_password = server_password or self.server_password self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.") + return True def connect(self) -> bool: """Connect the native client connection.""" diff --git a/tests/integration_tests/game_layer/observations/actions/__init__.py b/tests/integration_tests/game_layer/observations/actions/__init__.py new file mode 100644 index 00000000..be6c00e7 --- /dev/null +++ b/tests/integration_tests/game_layer/observations/actions/__init__.py @@ -0,0 +1 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK diff --git a/tests/integration_tests/game_layer/observations/actions/test_configure_actions.py b/tests/integration_tests/game_layer/observations/actions/test_configure_actions.py new file mode 100644 index 00000000..17e262d1 --- /dev/null +++ b/tests/integration_tests/game_layer/observations/actions/test_configure_actions.py @@ -0,0 +1,85 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +from ipaddress import IPv4Address + +from primaite.game.agent.actions import ConfigureDatabaseClientAction +from primaite.simulator.system.applications.database_client import DatabaseClient +from tests.conftest import ControlledAgent + + +class TestConfigureDatabaseAction: + def test_configure_ip_password(self, game_and_agent): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) + + # make sure there is a database client on this node + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(DatabaseClient) + db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] + + action = ( + "CONFIGURE_DATABASE_CLIENT", + { + "node_id": 0, + "options": { + "server_ip_address": "192.168.1.99", + "server_password": "admin123", + }, + }, + ) + agent.store_action(action) + game.step() + + assert db_client.server_ip_address == IPv4Address("192.168.1.99") + assert db_client.server_password == "admin123" + + def test_configure_ip(self, game_and_agent): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) + + # make sure there is a database client on this node + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(DatabaseClient) + db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] + + action = ( + "CONFIGURE_DATABASE_CLIENT", + { + "node_id": 0, + "options": { + "server_ip_address": "192.168.1.99", + }, + }, + ) + agent.store_action(action) + game.step() + + assert db_client.server_ip_address == IPv4Address("192.168.1.99") + assert db_client.server_password is None + + def test_configure_password(self, game_and_agent): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) + + # make sure there is a database client on this node + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(DatabaseClient) + db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] + old_ip = db_client.server_ip_address + + action = ( + "CONFIGURE_DATABASE_CLIENT", + { + "node_id": 0, + "options": { + "server_password": "admin123", + }, + }, + ) + agent.store_action(action) + game.step() + + assert db_client.server_ip_address == old_ip + assert db_client.server_password is "admin123" From ab73ac20e84766b5c77688b9adf0e02a98b22d6b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 1 Jul 2024 14:41:41 +0100 Subject: [PATCH 2/6] #2700 add ransomware configure action --- src/primaite/game/agent/actions.py | 27 ++++++- .../system/applications/database_client.py | 2 +- .../red_applications/ransomware_script.py | 14 +++- .../{observations => }/actions/__init__.py | 0 .../actions/test_configure_actions.py | 75 ++++++++++++++++++- 5 files changed, 113 insertions(+), 5 deletions(-) rename tests/integration_tests/game_layer/{observations => }/actions/__init__.py (100%) rename tests/integration_tests/game_layer/{observations => }/actions/test_configure_actions.py (52%) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 4cb31d25..9f2693e5 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -14,7 +14,7 @@ from abc import ABC, abstractmethod from typing import Dict, List, Literal, Optional, Tuple, TYPE_CHECKING, Union from gymnasium import spaces -from pydantic import BaseModel, Field, field_validator, ValidationInfo +from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo from primaite import getLogger from primaite.interface.request import RequestFormat @@ -252,6 +252,7 @@ class ConfigureDatabaseClientAction(AbstractAction): class _Opts(BaseModel): """Schema for options that can be passed to this action.""" + model_config = ConfigDict(extra="forbid") server_ip_address: Optional[str] = None server_password: Optional[str] = None @@ -267,6 +268,29 @@ class ConfigureDatabaseClientAction(AbstractAction): return ["network", "node", node_name, "application", "DatabaseClient", "configure", options] +class ConfigureRansomwareScriptAction(AbstractAction): + """Action which sets config parameters for a ransomware script on a node.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this option.""" + + model_config = ConfigDict(extra="forbid") + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + payload: Optional[str] = None + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, options: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + ConfigureRansomwareScriptAction._Opts.model_validate(options) # check that options adhere to schema + return ["network", "node", node_name, "application", "RansomwareScript", "configure", options] + + class NodeApplicationRemoveAction(AbstractAction): """Action which removes/uninstalls an application.""" @@ -1068,6 +1092,7 @@ class ActionManager: "NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction, "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction, "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, + "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, } """Dictionary which maps action type strings to the corresponding action class.""" diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index 6396c678..fcfd603b 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -104,7 +104,7 @@ class DatabaseClient(Application): success = self.configure(server_ip_address=ip, server_password=pw) return RequestResponse.from_bool(success) - rm.add_request("configure", RequestType(func=lambda request, context: _configure(request, context))) + rm.add_request("configure", RequestType(func=_configure)) return rm def execute(self) -> bool: 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 af4a59d4..46e42fc2 100644 --- a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py +++ b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py @@ -2,7 +2,7 @@ from ipaddress import IPv4Address from typing import Dict, Optional -from primaite.interface.request import RequestResponse +from primaite.interface.request import RequestFormat, 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 @@ -62,6 +62,15 @@ class RansomwareScript(Application): name="execute", request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.attack())), ) + + def _configure(request: RequestFormat, context: Dict) -> RequestResponse: + ip = request[-1].get("server_ip_address") + ip = None if ip is None else IPv4Address(ip) + pw = request[-1].get("server_password") + payload = request[-1].get("payload") + return RequestResponse.from_bool(self.configure(ip, pw, payload)) + + rm.add_request("configure", request_type=RequestType(func=_configure)) return rm def run(self) -> bool: @@ -91,7 +100,7 @@ class RansomwareScript(Application): server_ip_address: IPv4Address, server_password: Optional[str] = None, payload: Optional[str] = None, - ): + ) -> bool: """ Configure the Ransomware Script to communicate with a DatabaseService. @@ -108,6 +117,7 @@ class RansomwareScript(Application): self.sys_log.info( f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}." ) + return True def attack(self) -> bool: """Perform the attack steps after opening the application.""" diff --git a/tests/integration_tests/game_layer/observations/actions/__init__.py b/tests/integration_tests/game_layer/actions/__init__.py similarity index 100% rename from tests/integration_tests/game_layer/observations/actions/__init__.py rename to tests/integration_tests/game_layer/actions/__init__.py diff --git a/tests/integration_tests/game_layer/observations/actions/test_configure_actions.py b/tests/integration_tests/game_layer/actions/test_configure_actions.py similarity index 52% rename from tests/integration_tests/game_layer/observations/actions/test_configure_actions.py rename to tests/integration_tests/game_layer/actions/test_configure_actions.py index 17e262d1..5439f3c9 100644 --- a/tests/integration_tests/game_layer/observations/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -1,8 +1,12 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from ipaddress import IPv4Address -from primaite.game.agent.actions import ConfigureDatabaseClientAction +import pytest +from pydantic import ValidationError + +from primaite.game.agent.actions import ConfigureDatabaseClientAction, ConfigureRansomwareScriptAction from primaite.simulator.system.applications.database_client import DatabaseClient +from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript from tests.conftest import ControlledAgent @@ -83,3 +87,72 @@ class TestConfigureDatabaseAction: assert db_client.server_ip_address == old_ip assert db_client.server_password is "admin123" + + +class TestConfigureRansomwareScriptAction: + @pytest.mark.parametrize( + "options", + [ + {}, + {"server_ip_address": "181.181.181.181"}, + {"server_password": "admin123"}, + {"payload": "ENCRYPT"}, + { + "server_ip_address": "181.181.181.181", + "server_password": "admin123", + "payload": "ENCRYPT", + }, + ], + ) + def test_configure_ip_password(self, game_and_agent, options): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( + agent.action_manager + ) + + # make sure there is a database client on this node + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(RansomwareScript) + ransomware_script: RansomwareScript = client_1.software_manager.software["RansomwareScript"] + + old_ip = ransomware_script.server_ip_address + old_pw = ransomware_script.server_password + old_payload = ransomware_script.payload + + action = ( + "CONFIGURE_RANSOMWARE_SCRIPT", + {"node_id": 0, "options": options}, + ) + agent.store_action(action) + game.step() + + expected_ip = old_ip if "server_ip_address" not in options else IPv4Address(options["server_ip_address"]) + expected_pw = old_pw if "server_password" not in options else options["server_password"] + expected_payload = old_payload if "payload" not in options else options["payload"] + + assert ransomware_script.server_ip_address == expected_ip + assert ransomware_script.server_password == expected_pw + assert ransomware_script.payload == expected_payload + + def test_invalid_options(self, game_and_agent): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( + agent.action_manager + ) + + # make sure there is a database client on this node + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(RansomwareScript) + ransomware_script: RansomwareScript = client_1.software_manager.software["RansomwareScript"] + action = ( + "CONFIGURE_RANSOMWARE_SCRIPT", + { + "node_id": 0, + "options": {"server_password": "admin123", "bad_option": 70}, + }, + ) + agent.store_action(action) + with pytest.raises(ValidationError): + game.step() From bf8ec6083331adfcd089a55a4b0b6fc60b0423b3 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 1 Jul 2024 15:25:20 +0100 Subject: [PATCH 3/6] #2700 Add configure dosbot action --- src/primaite/game/agent/actions.py | 28 ++++++++++++ .../applications/red_applications/dos_bot.py | 17 +++++-- .../actions/test_configure_actions.py | 45 ++++++++++++++++++- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 9f2693e5..1de5276c 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -291,6 +291,33 @@ class ConfigureRansomwareScriptAction(AbstractAction): return ["network", "node", node_name, "application", "RansomwareScript", "configure", options] +class ConfigureDoSBotAction(AbstractAction): + """Action which sets config parameters for a DoS bot on a node.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this option.""" + + model_config = ConfigDict(extra="forbid") + target_ip_address: Optional[str] = None + target_port: Optional[str] = None + payload: Optional[str] = None + repeat: Optional[bool] = None + port_scan_p_of_success: Optional[float] = None + dos_intensity: Optional[float] = None + max_sessions: Optional[int] = None + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, options: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + self._Opts.model_validate(options) # check that options adhere to schema + return ["network", "node", node_name, "application", "DoSBot", "configure", options] + + class NodeApplicationRemoveAction(AbstractAction): """Action which removes/uninstalls an application.""" @@ -1093,6 +1120,7 @@ class ActionManager: "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction, "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, + "CONFIGURE_DOSBOT": ConfigureDoSBotAction, } """Dictionary which maps action type strings to the corresponding action class.""" 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 65e34227..dccf45f5 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -1,11 +1,11 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from enum import IntEnum from ipaddress import IPv4Address -from typing import Optional +from typing import Dict, Optional from primaite import getLogger from primaite.game.science import simulate_trial -from primaite.interface.request import RequestResponse +from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.database_client import DatabaseClient @@ -71,6 +71,14 @@ class DoSBot(DatabaseClient): request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.run())), ) + def _configure(request: RequestFormat, context: Dict) -> RequestResponse: + if "target_ip_address" in request[-1]: + request[-1]["target_ip_address"] = IPv4Address(request[-1]["target_ip_address"]) + if "target_port" in request[-1]: + request[-1]["target_port"] = Port[request[-1]["target_port"]] + return RequestResponse.from_bool(self.configure(**request[-1])) + + rm.add_request("configure", request_type=RequestType(func=_configure)) return rm def configure( @@ -82,7 +90,7 @@ class DoSBot(DatabaseClient): port_scan_p_of_success: float = 0.1, dos_intensity: float = 1.0, max_sessions: int = 1000, - ): + ) -> bool: """ Configure the Denial of Service bot. @@ -90,7 +98,7 @@ class DoSBot(DatabaseClient): :param: target_port: The port of the target service. Optional - Default is `Port.HTTP` :param: payload: The payload the DoS Bot will throw at the target service. Optional - Default is `None` :param: repeat: If True, the bot will maintain the attack. Optional - Default is `True` - :param: port_scan_p_of_success: The chance of the port scan being sucessful. Optional - Default is 0.1 (10%) + :param: port_scan_p_of_success: The chance of the port scan being successful. Optional - Default is 0.1 (10%) :param: dos_intensity: The intensity of the DoS attack. Multiplied with the application's max session - Default is 1.0 :param: max_sessions: The maximum number of sessions the DoS bot will attack with. Optional - Default is 1000 @@ -106,6 +114,7 @@ class DoSBot(DatabaseClient): f"{self.name}: Configured the {self.name} with {target_ip_address=}, {target_port=}, {payload=}, " f"{repeat=}, {port_scan_p_of_success=}, {dos_intensity=}, {max_sessions=}." ) + return True def run(self) -> bool: """Run the Denial of Service Bot.""" diff --git a/tests/integration_tests/game_layer/actions/test_configure_actions.py b/tests/integration_tests/game_layer/actions/test_configure_actions.py index 5439f3c9..6bcd3b52 100644 --- a/tests/integration_tests/game_layer/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -4,8 +4,14 @@ from ipaddress import IPv4Address import pytest from pydantic import ValidationError -from primaite.game.agent.actions import ConfigureDatabaseClientAction, ConfigureRansomwareScriptAction +from primaite.game.agent.actions import ( + ConfigureDatabaseClientAction, + ConfigureDoSBotAction, + ConfigureRansomwareScriptAction, +) +from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.database_client import DatabaseClient +from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript from tests.conftest import ControlledAgent @@ -156,3 +162,40 @@ class TestConfigureRansomwareScriptAction: agent.store_action(action) with pytest.raises(ValidationError): game.step() + + +class TestConfigureDoSBot: + def test_configure_DoSBot(self, game_and_agent): + game, agent = game_and_agent + agent: ControlledAgent + agent.action_manager.actions["CONFIGURE_DOSBOT"] = ConfigureDoSBotAction(agent.action_manager) + + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.software_manager.install(DoSBot) + dos_bot: DoSBot = client_1.software_manager.software["DoSBot"] + + action = ( + "CONFIGURE_DOSBOT", + { + "node_id": 0, + "options": { + "target_ip_address": "192.168.1.99", + "target_port": "POSTGRES_SERVER", + "payload": "HACC", + "repeat": False, + "port_scan_p_of_success": 0.875, + "dos_intensity": 0.75, + "max_sessions": 50, + }, + }, + ) + agent.store_action(action) + game.step() + + assert dos_bot.target_ip_address == IPv4Address("192.168.1.99") + assert dos_bot.target_port == Port.POSTGRES_SERVER + assert dos_bot.payload == "HACC" + assert not dos_bot.repeat + assert dos_bot.port_scan_p_of_success == 0.875 + assert dos_bot.dos_intensity == 0.75 + assert dos_bot.max_sessions == 50 From dc2c64b2f67fae73767ea0637d46fccb230ef1c1 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 1 Jul 2024 16:23:10 +0100 Subject: [PATCH 4/6] #2701 - Remove ip address option from node application install --- src/primaite/game/agent/actions.py | 3 +- .../game/agent/scripted_agents/tap001.py | 1 - .../simulator/network/hardware/base.py | 164 ++++++------------ .../configs/test_application_install.yaml | 1 - .../game_layer/test_actions.py | 2 +- 5 files changed, 53 insertions(+), 118 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index e165c9ad..3a21a95f 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -228,7 +228,7 @@ class NodeApplicationInstallAction(AbstractAction): super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes} - def form_request(self, node_id: int, application_name: str, ip_address: str) -> List[str]: + def form_request(self, node_id: int, application_name: str) -> List[str]: """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_name = self.manager.get_node_name_by_idx(node_id) if node_name is None: @@ -241,7 +241,6 @@ class NodeApplicationInstallAction(AbstractAction): "application", "install", application_name, - ip_address, ] diff --git a/src/primaite/game/agent/scripted_agents/tap001.py b/src/primaite/game/agent/scripted_agents/tap001.py index b1a378ef..c4f6062a 100644 --- a/src/primaite/game/agent/scripted_agents/tap001.py +++ b/src/primaite/game/agent/scripted_agents/tap001.py @@ -55,7 +55,6 @@ class TAP001(AbstractScriptedAgent): return "NODE_APPLICATION_INSTALL", { "node_id": self.starting_node_idx, "application_name": "RansomwareScript", - "ip_address": self.ip_address, } return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0} diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 01745215..1982b08f 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -6,7 +6,7 @@ import secrets from abc import ABC, abstractmethod from ipaddress import IPv4Address, IPv4Network from pathlib import Path -from typing import Any, Dict, Optional, Type, TypeVar, Union +from typing import Any, Dict, Optional, TypeVar, Union from prettytable import MARKDOWN, PrettyTable from pydantic import BaseModel, Field @@ -884,6 +884,54 @@ class Node(SimComponent): More information in user guide and docstring for SimComponent._init_request_manager. """ + + def _install_application(request: RequestFormat, context: Dict) -> RequestResponse: + """ + Allows agents to install applications to the node. + + :param request: list containing the application name as the only element + :type application: str + """ + application_name = request[0] + if self.software_manager.software.get(application_name): + self.sys_log.warning(f"Can't install {application_name}. It's already installed.") + return RequestResponse.from_bool(False) + application_class = Application._application_registry[application_name] + self.software_manager.install(application_class) + application_instance = self.software_manager.software.get(application_name) + self.applications[application_instance.uuid] = application_instance + _LOGGER.debug(f"Added application {application_instance.name} to node {self.hostname}") + self._application_request_manager.add_request( + application_name, RequestType(func=application_instance._request_manager) + ) + application_instance.install() + if application_name in self.software_manager.software: + return RequestResponse.from_bool(True) + else: + return RequestResponse.from_bool(False) + + def _uninstall_application(request: RequestFormat, context: Dict) -> RequestResponse: + """ + 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. + """ + application_name = request[0] + if application_name not in self.software_manager.software: + self.sys_log.warning(f"Can't uninstall {application_name}. It's not installed.") + return RequestResponse.from_bool(False) + + application_instance = self.software_manager.software.get(application_name) + self.software_manager.uninstall(application_instance.name) + if application_instance.name not in self.software_manager.software: + return RequestResponse.from_bool(True) + else: + return RequestResponse.from_bool(False) + _node_is_on = Node._NodeIsOnValidator(node=self) rm = super()._init_request_manager() @@ -940,25 +988,8 @@ class Node(SimComponent): 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])) - ) - ), - ) + self._application_manager.add_request(name="install", request_type=RequestType(func=_install_application)) + self._application_manager.add_request(name="uninstall", request_type=RequestType(func=_uninstall_application)) return rm @@ -966,29 +997,6 @@ class Node(SimComponent): """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. @@ -1417,76 +1425,6 @@ class Node(SimComponent): self.sys_log.info(f"Uninstalled application {application.name}") 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 - _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 - application_instance.install() - 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 diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index 87402f73..a4e898ae 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/tests/assets/configs/test_application_install.yaml @@ -683,7 +683,6 @@ agents: options: node_id: 0 application_name: DoSBot - ip_address: 192.168.1.14 79: action: NODE_APPLICATION_REMOVE options: diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 0dcf125d..a1005f34 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -557,7 +557,7 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl assert client_1.software_manager.software.get("DoSBot") is None - action = ("NODE_APPLICATION_INSTALL", {"node_id": 0, "application_name": "DoSBot", "ip_address": "192.168.1.14"}) + action = ("NODE_APPLICATION_INSTALL", {"node_id": 0, "application_name": "DoSBot"}) agent.store_action(action) game.step() From b27ac52d9ecdb05268175413b283bec4e9101946 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 2 Jul 2024 11:10:19 +0100 Subject: [PATCH 5/6] #2700 add E2E tests for application configure actions --- src/primaite/game/agent/actions.py | 18 +-- src/primaite/game/game.py | 3 +- .../applications/red_applications/dos_bot.py | 2 +- .../configs/install_and_configure_apps.yaml | 142 ++++++++++++++++++ .../configs/test_application_install.yaml | 9 ++ .../test_uc2_data_manipulation_scenario.py | 1 + .../actions/test_configure_actions.py | 115 ++++++++++++-- 7 files changed, 267 insertions(+), 23 deletions(-) create mode 100644 tests/assets/configs/install_and_configure_apps.yaml diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 60ff19e5..b3b7189c 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -258,13 +258,13 @@ class ConfigureDatabaseClientAction(AbstractAction): def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) - def form_request(self, node_id: int, options: Dict) -> RequestFormat: + def form_request(self, node_id: int, config: Dict) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" node_name = self.manager.get_node_name_by_idx(node_id) if node_name is None: return ["do_nothing"] - ConfigureDatabaseClientAction._Opts.model_validate(options) # check that options adhere to schema - return ["network", "node", node_name, "application", "DatabaseClient", "configure", options] + ConfigureDatabaseClientAction._Opts.model_validate(config) # check that options adhere to schema + return ["network", "node", node_name, "application", "DatabaseClient", "configure", config] class ConfigureRansomwareScriptAction(AbstractAction): @@ -281,13 +281,13 @@ class ConfigureRansomwareScriptAction(AbstractAction): def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) - def form_request(self, node_id: int, options: Dict) -> RequestFormat: + def form_request(self, node_id: int, config: Dict) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" node_name = self.manager.get_node_name_by_idx(node_id) if node_name is None: return ["do_nothing"] - ConfigureRansomwareScriptAction._Opts.model_validate(options) # check that options adhere to schema - return ["network", "node", node_name, "application", "RansomwareScript", "configure", options] + ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema + return ["network", "node", node_name, "application", "RansomwareScript", "configure", config] class ConfigureDoSBotAction(AbstractAction): @@ -308,13 +308,13 @@ class ConfigureDoSBotAction(AbstractAction): def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) - def form_request(self, node_id: int, options: Dict) -> RequestFormat: + def form_request(self, node_id: int, config: Dict) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" node_name = self.manager.get_node_name_by_idx(node_id) if node_name is None: return ["do_nothing"] - self._Opts.model_validate(options) # check that options adhere to schema - return ["network", "node", node_name, "application", "DoSBot", "configure", options] + self._Opts.model_validate(config) # check that options adhere to schema + return ["network", "node", node_name, "application", "DoSBot", "configure", config] class NodeApplicationRemoveAction(AbstractAction): diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 05210278..89102afb 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -313,7 +313,8 @@ class PrimaiteGame: if "options" in service_cfg: opt = service_cfg["options"] new_service.password = opt.get("db_password", None) - new_service.configure_backup(backup_server=IPv4Address(opt.get("backup_server_ip"))) + if "backup_server_ip" in opt: + new_service.configure_backup(backup_server=IPv4Address(opt.get("backup_server_ip"))) if service_type == "FTPServer": if "options" in service_cfg: opt = service_cfg["options"] 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 17478b71..01a375ee 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -135,7 +135,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"): self.sys_log.warning( f"{self.name} is not properly configured. {self.target_ip_address=}, {self.target_port=}" ) - return True + return False self.clear_connections() self._perform_port_scan(p_of_success=self.port_scan_p_of_success) diff --git a/tests/assets/configs/install_and_configure_apps.yaml b/tests/assets/configs/install_and_configure_apps.yaml new file mode 100644 index 00000000..6b548f7e --- /dev/null +++ b/tests/assets/configs/install_and_configure_apps.yaml @@ -0,0 +1,142 @@ +io_settings: + save_step_metadata: false + save_pcap_logs: false + save_sys_logs: false + save_agent_actions: false + +game: + max_episode_length: 256 + ports: + - ARP + - DNS + protocols: + - ICMP + - TCP + +agents: + - ref: agent_1 + team: BLUE + type: ProxyAgent + + observation_space: null + action_space: + action_list: + - type: DONOTHING + - type: NODE_APPLICATION_INSTALL + - type: CONFIGURE_DATABASE_CLIENT + - type: CONFIGURE_DOSBOT + - type: CONFIGURE_RANSOMWARE_SCRIPT + - type: NODE_APPLICATION_REMOVE + action_map: + 0: + action: DONOTHING + options: {} + 1: + action: NODE_APPLICATION_INSTALL + options: + node_id: 0 + application_name: DatabaseClient + 2: + action: NODE_APPLICATION_INSTALL + options: + node_id: 1 + application_name: RansomwareScript + 3: + action: NODE_APPLICATION_INSTALL + options: + node_id: 2 + application_name: DoSBot + 4: + action: CONFIGURE_DATABASE_CLIENT + options: + node_id: 0 + config: + server_ip_address: 10.0.0.5 + 5: + action: CONFIGURE_DATABASE_CLIENT + options: + node_id: 0 + config: + server_password: correct_password + 6: + action: CONFIGURE_RANSOMWARE_SCRIPT + options: + node_id: 1 + config: + server_ip_address: 10.0.0.5 + server_password: correct_password + payload: ENCRYPT + 7: + action: CONFIGURE_DOSBOT + options: + node_id: 2 + config: + target_ip_address: 10.0.0.5 + target_port: POSTGRES_SERVER + payload: DELETE + repeat: true + port_scan_p_of_success: 1.0 + dos_intensity: 1.0 + max_sessions: 1000 + 8: + action: NODE_APPLICATION_INSTALL + options: + node_id: 1 + application_name: DatabaseClient + options: + nodes: + - node_name: client_1 + - node_name: client_2 + - node_name: client_3 + ip_list: [] + reward_function: + reward_components: + - type: DUMMY + +simulation: + network: + nodes: + - type: computer + hostname: client_1 + ip_address: 10.0.0.2 + subnet_mask: 255.255.255.0 + default_gateway: 10.0.0.1 + - type: computer + hostname: client_2 + ip_address: 10.0.0.3 + subnet_mask: 255.255.255.0 + default_gateway: 10.0.0.1 + - type: computer + hostname: client_3 + ip_address: 10.0.0.4 + subnet_mask: 255.255.255.0 + default_gateway: 10.0.0.1 + - type: switch + hostname: switch_1 + num_ports: 8 + - type: server + hostname: server_1 + ip_address: 10.0.0.5 + subnet_mask: 255.255.255.0 + default_gateway: 10.0.0.1 + services: + - type: DatabaseService + options: + db_password: correct_password + links: + - endpoint_a_hostname: client_1 + endpoint_a_port: 1 + endpoint_b_hostname: switch_1 + endpoint_b_port: 1 + - endpoint_a_hostname: client_2 + endpoint_a_port: 1 + endpoint_b_hostname: switch_1 + endpoint_b_port: 2 + - endpoint_a_hostname: client_3 + endpoint_a_port: 1 + endpoint_b_hostname: switch_1 + endpoint_b_port: 3 + - endpoint_a_hostname: server_1 + endpoint_a_port: 1 + endpoint_b_hostname: switch_1 + endpoint_b_port: 8 diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index a4e898ae..3a3a6890 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/tests/assets/configs/test_application_install.yaml @@ -260,6 +260,7 @@ agents: - type: NODE_APPLICATION_INSTALL - type: NODE_APPLICATION_REMOVE - type: NODE_APPLICATION_EXECUTE + - type: CONFIGURE_DOSBOT action_map: 0: @@ -698,6 +699,14 @@ agents: options: node_id: 0 application_id: 0 + 82: + action: CONFIGURE_DOSBOT + options: + node_id: 0 + config: + target_ip_address: 192.168.1.14 + target_port: POSTGRES_SERVER + 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 e6cd113f..7ec38d72 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -69,6 +69,7 @@ def test_application_install_uninstall_on_uc2(): env.step(0) # Test we can now execute the DoSBot app + env.step(82) # configure dos bot with ip address and port _, _, _, _, info = env.step(81) assert info["agent_actions"]["defender"].response.status == "success" diff --git a/tests/integration_tests/game_layer/actions/test_configure_actions.py b/tests/integration_tests/game_layer/actions/test_configure_actions.py index 6bcd3b52..b7acc8a8 100644 --- a/tests/integration_tests/game_layer/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -9,12 +9,19 @@ from primaite.game.agent.actions import ( ConfigureDoSBotAction, ConfigureRansomwareScriptAction, ) +from primaite.session.environment import PrimaiteGymEnv +from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.applications.application import ApplicationOperatingState from primaite.simulator.system.applications.database_client import DatabaseClient from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript +from primaite.simulator.system.services.database.database_service import DatabaseService +from tests import TEST_ASSETS_ROOT from tests.conftest import ControlledAgent +APP_CONFIG_YAML = TEST_ASSETS_ROOT / "configs/install_and_configure_apps.yaml" + class TestConfigureDatabaseAction: def test_configure_ip_password(self, game_and_agent): @@ -31,7 +38,7 @@ class TestConfigureDatabaseAction: "CONFIGURE_DATABASE_CLIENT", { "node_id": 0, - "options": { + "config": { "server_ip_address": "192.168.1.99", "server_password": "admin123", }, @@ -57,7 +64,7 @@ class TestConfigureDatabaseAction: "CONFIGURE_DATABASE_CLIENT", { "node_id": 0, - "options": { + "config": { "server_ip_address": "192.168.1.99", }, }, @@ -83,7 +90,7 @@ class TestConfigureDatabaseAction: "CONFIGURE_DATABASE_CLIENT", { "node_id": 0, - "options": { + "config": { "server_password": "admin123", }, }, @@ -97,7 +104,7 @@ class TestConfigureDatabaseAction: class TestConfigureRansomwareScriptAction: @pytest.mark.parametrize( - "options", + "config", [ {}, {"server_ip_address": "181.181.181.181"}, @@ -110,7 +117,7 @@ class TestConfigureRansomwareScriptAction: }, ], ) - def test_configure_ip_password(self, game_and_agent, options): + def test_configure_ip_password(self, game_and_agent, config): game, agent = game_and_agent agent: ControlledAgent agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( @@ -128,20 +135,20 @@ class TestConfigureRansomwareScriptAction: action = ( "CONFIGURE_RANSOMWARE_SCRIPT", - {"node_id": 0, "options": options}, + {"node_id": 0, "config": config}, ) agent.store_action(action) game.step() - expected_ip = old_ip if "server_ip_address" not in options else IPv4Address(options["server_ip_address"]) - expected_pw = old_pw if "server_password" not in options else options["server_password"] - expected_payload = old_payload if "payload" not in options else options["payload"] + expected_ip = old_ip if "server_ip_address" not in config else IPv4Address(config["server_ip_address"]) + expected_pw = old_pw if "server_password" not in config else config["server_password"] + expected_payload = old_payload if "payload" not in config else config["payload"] assert ransomware_script.server_ip_address == expected_ip assert ransomware_script.server_password == expected_pw assert ransomware_script.payload == expected_payload - def test_invalid_options(self, game_and_agent): + def test_invalid_config(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( @@ -156,7 +163,7 @@ class TestConfigureRansomwareScriptAction: "CONFIGURE_RANSOMWARE_SCRIPT", { "node_id": 0, - "options": {"server_password": "admin123", "bad_option": 70}, + "config": {"server_password": "admin123", "bad_option": 70}, }, ) agent.store_action(action) @@ -178,7 +185,7 @@ class TestConfigureDoSBot: "CONFIGURE_DOSBOT", { "node_id": 0, - "options": { + "config": { "target_ip_address": "192.168.1.99", "target_port": "POSTGRES_SERVER", "payload": "HACC", @@ -199,3 +206,87 @@ class TestConfigureDoSBot: assert dos_bot.port_scan_p_of_success == 0.875 assert dos_bot.dos_intensity == 0.75 assert dos_bot.max_sessions == 50 + + +class TestConfigureYAML: + def test_configure_db_client(self): + env = PrimaiteGymEnv(env_config=APP_CONFIG_YAML) + + # make sure there's no db client on the node yet + client_1 = env.game.simulation.network.get_node_by_hostname("client_1") + assert client_1.software_manager.software.get("DatabaseClient") is None + + # take the install action, check that the db gets installed, step to get it to finish installing + env.step(1) + db_client: DatabaseClient = client_1.software_manager.software.get("DatabaseClient") + assert isinstance(db_client, DatabaseClient) + assert db_client.operating_state == ApplicationOperatingState.INSTALLING + env.step(0) + env.step(0) + env.step(0) + env.step(0) + + # configure the ip address and check that it changes, but password doesn't change + assert db_client.server_ip_address is None + assert db_client.server_password is None + env.step(4) + assert db_client.server_ip_address == IPv4Address("10.0.0.5") + assert db_client.server_password is None + + # configure the password and check that it changes, make sure this lets us connect to the db + assert not db_client.connect() + env.step(5) + assert db_client.server_password == "correct_password" + assert db_client.connect() + + def test_configure_ransomware_script(self): + env = PrimaiteGymEnv(env_config=APP_CONFIG_YAML) + client_2 = env.game.simulation.network.get_node_by_hostname("client_2") + assert client_2.software_manager.software.get("RansomwareScript") is None + + # install ransomware script + env.step(2) + ransom = client_2.software_manager.software.get("RansomwareScript") + assert isinstance(ransom, RansomwareScript) + assert ransom.operating_state == ApplicationOperatingState.INSTALLING + env.step(0) + env.step(0) + env.step(0) + env.step(0) + + # make sure it's not working yet because it's not configured and there's no db client + assert not ransom.attack() + env.step(8) # install db client on the same node + env.step(0) + env.step(0) + env.step(0) + env.step(0) # let it finish installing + assert not ransom.attack() + + # finally, configure the ransomware script with ip and password + env.step(6) + assert ransom.attack() + + db_server = env.game.simulation.network.get_node_by_hostname("server_1") + db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService") + assert db_service.db_file.health_status == FileSystemItemHealthStatus.CORRUPT + + def test_configure_dos_bot(self): + env = PrimaiteGymEnv(env_config=APP_CONFIG_YAML) + client_3 = env.game.simulation.network.get_node_by_hostname("client_3") + assert client_3.software_manager.software.get("DoSBot") is None + + # install DoSBot + env.step(3) + bot = client_3.software_manager.software.get("DoSBot") + assert isinstance(bot, DoSBot) + assert bot.operating_state == ApplicationOperatingState.INSTALLING + env.step(0) + env.step(0) + env.step(0) + env.step(0) + + # make sure dos bot doesn't work before being configured + assert not bot.run() + env.step(7) + assert bot.run() From feabe5117c5d1471e61aeefc1f7609b42c819b28 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 2 Jul 2024 12:48:23 +0100 Subject: [PATCH 6/6] #2700 Fix docstrings in application configure methods --- src/primaite/simulator/network/hardware/base.py | 15 +++++++++++---- .../applications/red_applications/dos_bot.py | 15 +++++++++++++++ .../red_applications/ransomware_script.py | 10 ++++++++++ 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 1982b08f..6942d280 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -890,7 +890,11 @@ class Node(SimComponent): Allows agents to install applications to the node. :param request: list containing the application name as the only element - :type application: str + :type request: RequestFormat + :param context: additional context for resolving this action, currently unused + :type context: dict + :return: Request response with a success code if the application was installed. + :rtype: RequestResponse """ application_name = request[0] if self.software_manager.software.get(application_name): @@ -916,9 +920,12 @@ class Node(SimComponent): 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. + :param request: list containing the application name as the only element + :type request: RequestFormat + :param context: additional context for resolving this action, currently unused + :type context: dict + :return: Request response with a success code if the application was uninstalled. + :rtype: RequestResponse """ application_name = request[0] if application_name not in self.software_manager.software: 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 01a375ee..fcad3b3e 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -72,6 +72,16 @@ class DoSBot(DatabaseClient, identifier="DoSBot"): ) def _configure(request: RequestFormat, context: Dict) -> RequestResponse: + """ + Configure the DoSBot. + + :param request: List with one element that is a dict of options to pass to the configure method. + :type request: RequestFormat + :param context: additional context for resolving this action, currently unused + :type context: dict + :return: Request Response object with a success code determining if the configuration was successful. + :rtype: RequestResponse + """ if "target_ip_address" in request[-1]: request[-1]["target_ip_address"] = IPv4Address(request[-1]["target_ip_address"]) if "target_port" in request[-1]: @@ -102,6 +112,8 @@ class DoSBot(DatabaseClient, identifier="DoSBot"): :param: dos_intensity: The intensity of the DoS attack. Multiplied with the application's max session - Default is 1.0 :param: max_sessions: The maximum number of sessions the DoS bot will attack with. Optional - Default is 1000 + :return: Always returns True + :rtype: bool """ self.target_ip_address = target_ip_address self.target_port = target_port @@ -126,6 +138,9 @@ class DoSBot(DatabaseClient, identifier="DoSBot"): The main application loop for the Denial of Service bot. The loop goes through the stages of a DoS attack. + + :return: True if the application loop could be executed, False otherwise. + :rtype: bool """ if not self._can_perform_action(): return False 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 8d9d0d18..71e422c3 100644 --- a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py +++ b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py @@ -64,6 +64,16 @@ class RansomwareScript(Application, identifier="RansomwareScript"): ) def _configure(request: RequestFormat, context: Dict) -> RequestResponse: + """ + Request for configuring the target database and payload. + + :param request: Request with one element contianing a dict of parameters for the configure method. + :type request: RequestFormat + :param context: additional context for resolving this action, currently unused + :type context: dict + :return: RequestResponse object with a success code reflecting whether the configuration could be applied. + :rtype: RequestResponse + """ ip = request[-1].get("server_ip_address") ip = None if ip is None else IPv4Address(ip) pw = request[-1].get("server_password")