From d8154bbebd4e6d98aaf6cf51628be1b176ad00b8 Mon Sep 17 00:00:00 2001 From: Jake Walker Date: Tue, 21 Nov 2023 11:43:47 +0000 Subject: [PATCH] Add tests for data manipulation bot attack stages --- .../test_data_manipulation_bot.py | 61 +++++++++++++++++-- 1 file changed, 57 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py b/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py index dd785cc1..5127254c 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py @@ -1,20 +1,73 @@ from ipaddress import IPv4Address +import pytest + from primaite.simulator.network.hardware.base import Node from primaite.simulator.network.networks import arcd_uc2_network from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port -from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot +from primaite.simulator.system.services.red_services.data_manipulation_bot import ( + DataManipulationAttackStage, + DataManipulationBot, +) -def test_creation(): +@pytest.fixture(scope="function") +def dm_client() -> Node: network = arcd_uc2_network() + return network.get_node_by_hostname("client_1") - client_1: Node = network.get_node_by_hostname("client_1") - data_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"] +@pytest.fixture +def dm_bot(dm_client) -> DataManipulationBot: + return dm_client.software_manager.software["DataManipulationBot"] + + +def test_create_dm_bot(dm_client): + data_manipulation_bot: DataManipulationBot = dm_client.software_manager.software["DataManipulationBot"] assert data_manipulation_bot.name == "DataManipulationBot" assert data_manipulation_bot.port == Port.POSTGRES_SERVER assert data_manipulation_bot.protocol == IPProtocol.TCP assert data_manipulation_bot.payload == "DROP TABLE IF EXISTS user;" + + +def test_dm_bot_logon(dm_bot): + dm_bot.attack_stage = DataManipulationAttackStage.NOT_STARTED + + dm_bot._logon() + + assert dm_bot.attack_stage == DataManipulationAttackStage.LOGON + + +def test_dm_bot_perform_port_scan_no_success(dm_bot): + dm_bot.attack_stage = DataManipulationAttackStage.LOGON + + dm_bot._perform_port_scan(p_of_success=0.0) + + assert dm_bot.attack_stage == DataManipulationAttackStage.LOGON + + +def test_dm_bot_perform_port_scan_success(dm_bot): + dm_bot.attack_stage = DataManipulationAttackStage.LOGON + + dm_bot._perform_port_scan(p_of_success=1.0) + + assert dm_bot.attack_stage == DataManipulationAttackStage.PORT_SCAN + + +def test_dm_bot_perform_data_manipulation_no_success(dm_bot): + dm_bot.attack_stage = DataManipulationAttackStage.PORT_SCAN + + dm_bot._perform_data_manipulation(p_of_success=0.0) + + assert dm_bot.attack_stage == DataManipulationAttackStage.PORT_SCAN + + +def test_dm_bot_perform_data_manipulation_success(dm_bot): + dm_bot.attack_stage = DataManipulationAttackStage.PORT_SCAN + + dm_bot._perform_data_manipulation(p_of_success=1.0) + + assert dm_bot.attack_stage in (DataManipulationAttackStage.COMPLETE, DataManipulationAttackStage.FAILED) + assert dm_bot.connected