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

#2700 add E2E tests for application configure actions

Browse Source
This commit is contained in:
Marek Wolan
2024-07-02 11:10:19 +01:00
parent 33b9c5f6b3
commit b27ac52d9e
7 changed files with 267 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/primaite/game/agent/actions.py
View File

@@ -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):

3
src/primaite/game/game.py
View File

@@ -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"]

2
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -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)

142
tests/assets/configs/install_and_configure_apps.yaml Normal file
View File

@@ -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

9
tests/assets/configs/test_application_install.yaml
View File

@@ -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

1
tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
View File

@@ -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"

115
tests/integration_tests/game_layer/actions/test_configure_actions.py
View File

@@ -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()