From b27ac52d9ecdb05268175413b283bec4e9101946 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 2 Jul 2024 11:10:19 +0100 Subject: [PATCH] #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()