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

Merge 'origin/dev' into feature/2319-database-admin

Browse Source
This commit is contained in:
Marek Wolan
2024-02-29 13:21:44 +00:00
parent af8ca82fcb 312b5c35bd
commit f7c535b557
62 changed files with 865 additions and 1163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/assets/configs/bad_primaite_session.yaml
View File

@@ -589,15 +589,16 @@ simulation:
hostname: web_server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.10
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server

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

@@ -593,15 +593,16 @@ simulation:
hostname: web_server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.10
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server
@@ -624,7 +625,7 @@ simulation:
dns_server: 192.168.1.10
services:
- ref: backup_service
type: DatabaseBackup
type: FTPServer
- ref: security_suite
type: server

10
tests/assets/configs/multi_agent_session.yaml
View File

@@ -1043,16 +1043,16 @@ simulation:
hostname: web_server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.10
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server
type: server
@@ -1074,7 +1074,7 @@ simulation:
dns_server: 192.168.1.10
services:
- ref: backup_service
type: DatabaseBackup
type: FTPServer
- ref: security_suite
type: server

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

@@ -599,15 +599,16 @@ simulation:
hostname: web_server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.10
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server
@@ -630,7 +631,7 @@ simulation:
dns_server: 192.168.1.10
services:
- ref: backup_service
type: DatabaseBackup
type: FTPServer
- ref: security_suite
type: server

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

@@ -600,15 +600,16 @@ simulation:
hostname: web_server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.10
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server
@@ -631,7 +632,7 @@ simulation:
dns_server: 192.168.1.10
services:
- ref: backup_service
type: DatabaseBackup
type: FTPServer
- ref: security_suite
type: server

2
tests/conftest.py
View File

@@ -509,6 +509,6 @@ def game_and_agent():
reward_function=reward_function,
)
game.agents.append(test_agent)
game.agents["test_agent"] = test_agent
return (game, test_agent)

4
tests/e2e_integration_tests/environments/test_sb3_environment.py
View File

@@ -11,14 +11,12 @@ from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
# @pytest.mark.skip(reason="no way of currently testing this")
def test_sb3_compatibility():
"""Test that the Gymnasium environment can be used with an SB3 agent."""
with open(example_config_path(), "r") as f:
cfg = yaml.safe_load(f)
game = PrimaiteGame.from_config(cfg)
gym = PrimaiteGymEnv(game=game)
gym = PrimaiteGymEnv(game_config=cfg)
model = PPO("MlpPolicy", gym)
model.learn(total_timesteps=1000)

16
tests/integration_tests/game_configuration.py
View File

@@ -42,20 +42,20 @@ def test_example_config():
assert len(game.agents) == 4 # red, blue and 2 green agents
# green agent 1
assert game.agents[0].agent_name == "client_2_green_user"
assert isinstance(game.agents[0], RandomAgent)
assert "client_2_green_user" in game.agents
assert isinstance(game.agents["client_2_green_user"], RandomAgent)
# green agent 2
assert game.agents[1].agent_name == "client_1_green_user"
assert isinstance(game.agents[1], RandomAgent)
assert "client_1_green_user" in game.agents
assert isinstance(game.agents["client_1_green_user"], RandomAgent)
# red agent
assert game.agents[2].agent_name == "client_1_data_manipulation_red_bot"
assert isinstance(game.agents[2], DataManipulationAgent)
assert "client_1_data_manipulation_red_bot" in game.agents
assert isinstance(game.agents["client_1_data_manipulation_red_bot"], DataManipulationAgent)
# blue agent
assert game.agents[3].agent_name == "defender"
assert isinstance(game.agents[3], ProxyAgent)
assert "defender" in game.agents
assert isinstance(game.agents["defender"], ProxyAgent)
network: Network = game.simulation.network

120
tests/integration_tests/network/test_capture_nmne.py Normal file
View File

@@ -0,0 +1,120 @@
from primaite.game.agent.observations import NicObservation
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.nmne import set_nmne_config
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
def test_capture_nmne(uc2_network):
"""
Conducts a test to verify that Malicious Network Events (MNEs) are correctly captured.
This test involves a web server querying a database server and checks if the MNEs are captured
based on predefined keywords in the network configuration. Specifically, it checks the capture
of the "DELETE" SQL command as a malicious network event.
"""
web_server: Server = uc2_network.get_node_by_hostname("web_server") # noqa
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] # noqa
db_client.connect()
db_server: Server = uc2_network.get_node_by_hostname("database_server") # noqa
web_server_nic = web_server.network_interface[1]
db_server_nic = db_server.network_interface[1]
# Set the NMNE configuration to capture DELETE queries as MNEs
nmne_config = {
"capture_nmne": True, # Enable the capture of MNEs
"nmne_capture_keywords": ["DELETE"], # Specify "DELETE" SQL command as a keyword for MNE detection
}
# Apply the NMNE configuration settings
set_nmne_config(nmne_config)
# Assert that initially, there are no captured MNEs on both web and database servers
assert web_server_nic.describe_state()["nmne"] == {}
assert db_server_nic.describe_state()["nmne"] == {}
# Perform a "SELECT" query
db_client.query("SELECT")
# Check that it does not trigger an MNE capture.
assert web_server_nic.describe_state()["nmne"] == {}
assert db_server_nic.describe_state()["nmne"] == {}
# Perform a "DELETE" query
db_client.query("DELETE")
# Check that the web server's outbound interface and the database server's inbound interface register the MNE
assert web_server_nic.describe_state()["nmne"] == {"direction": {"outbound": {"keywords": {"*": 1}}}}
assert db_server_nic.describe_state()["nmne"] == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "SELECT" query
db_client.query("SELECT")
# Check that no additional MNEs are captured
assert web_server_nic.describe_state()["nmne"] == {"direction": {"outbound": {"keywords": {"*": 1}}}}
assert db_server_nic.describe_state()["nmne"] == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "DELETE" query
db_client.query("DELETE")
# Check that the web server and database server interfaces register an additional MNE
assert web_server_nic.describe_state()["nmne"] == {"direction": {"outbound": {"keywords": {"*": 2}}}}
assert db_server_nic.describe_state()["nmne"] == {"direction": {"inbound": {"keywords": {"*": 2}}}}
def test_capture_nmne_observations(uc2_network):
"""
Tests the NicObservation class's functionality within a simulated network environment.
This test ensures the observation space, as defined by instances of NicObservation, accurately reflects the
number of MNEs detected based on network activities over multiple iterations.
The test employs a series of "DELETE" SQL operations, considered as MNEs, to validate the dynamic update
and accuracy of the observation space related to network interface conditions. It confirms that the
observed NIC states match expected MNE activity levels.
"""
# Initialise a new Simulation instance and assign the test network to it.
sim = Simulation()
sim.network = uc2_network
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect()
# Set the NMNE configuration to capture DELETE queries as MNEs
nmne_config = {
"capture_nmne": True, # Enable the capture of MNEs
"nmne_capture_keywords": ["DELETE"], # Specify "DELETE" SQL command as a keyword for MNE detection
}
# Apply the NMNE configuration settings
set_nmne_config(nmne_config)
# Define observations for the NICs of the database and web servers
db_server_nic_obs = NicObservation(where=["network", "nodes", "database_server", "NICs", 1])
web_server_nic_obs = NicObservation(where=["network", "nodes", "web_server", "NICs", 1])
# Iterate through a set of test cases to simulate multiple DELETE queries
for i in range(1, 20):
# Perform a "DELETE" query each iteration
db_client.query("DELETE")
# Observe the current state of NMNEs from the NICs of both the database and web servers
db_nic_obs = db_server_nic_obs.observe(sim.describe_state())["nmne"]
web_nic_obs = web_server_nic_obs.observe(sim.describe_state())["nmne"]
# Define expected NMNE values based on the iteration count
if i > 10:
expected_nmne = 3 # High level of detected MNEs after 10 iterations
elif i > 5:
expected_nmne = 2 # Moderate level after more than 5 iterations
elif i > 0:
expected_nmne = 1 # Low level detected after just starting
else:
expected_nmne = 0 # No MNEs detected
# Assert that the observed NMNEs match the expected values for both NICs
assert web_nic_obs["outbound"] == expected_nmne
assert db_nic_obs["inbound"] == expected_nmne

21
tests/unit_tests/_primaite/_simulator/_domain/test_account.py
View File

@@ -7,26 +7,11 @@ from primaite.simulator.domain.account import Account, AccountType
@pytest.fixture(scope="function")
def account() -> Account:
acct = Account(username="Jake", password="totally_hashed_password", account_type=AccountType.USER)
acct.set_original_state()
return acct
def test_original_state(account):
"""Test the original state - see if it resets properly"""
account.log_on()
account.log_off()
account.disable()
state = account.describe_state()
assert state["num_logons"] is 1
assert state["num_logoffs"] is 1
assert state["num_group_changes"] is 0
assert state["username"] is "Jake"
assert state["password"] is "totally_hashed_password"
assert state["account_type"] is AccountType.USER.value
assert state["enabled"] is False
account.reset_component_for_episode(episode=1)
state = account.describe_state()
assert state["num_logons"] is 0
assert state["num_logoffs"] is 0
@@ -39,13 +24,7 @@ def test_original_state(account):
account.log_on()
account.log_off()
account.disable()
account.set_original_state()
account.log_on()
state = account.describe_state()
assert state["num_logons"] is 2
account.reset_component_for_episode(episode=2)
state = account.describe_state()
assert state["num_logons"] is 1
assert state["num_logoffs"] is 1

32
tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py
View File

@@ -185,38 +185,6 @@ def test_get_file(file_system):
file_system.show(full=True)
def test_reset_file_system(file_system):
# file and folder that existed originally
file_system.create_file(file_name="test_file.zip")
file_system.create_folder(folder_name="test_folder")
file_system.set_original_state()
# create a new file
file_system.create_file(file_name="new_file.txt")
# create a new folder
file_system.create_folder(folder_name="new_folder")
# delete the file that existed originally
file_system.delete_file(folder_name="root", file_name="test_file.zip")
assert file_system.get_file(folder_name="root", file_name="test_file.zip") is None
# delete the folder that existed originally
file_system.delete_folder(folder_name="test_folder")
assert file_system.get_folder(folder_name="test_folder") is None
# reset
file_system.reset_component_for_episode(episode=1)
# deleted original file and folder should be back
assert file_system.get_file(folder_name="root", file_name="test_file.zip")
assert file_system.get_folder(folder_name="test_folder")
# new file and folder should be removed
assert file_system.get_file(folder_name="root", file_name="new_file.txt") is None
assert file_system.get_folder(folder_name="new_folder") is None
@pytest.mark.skip(reason="Skipping until we tackle serialisation")
def test_serialisation(file_system):
"""Test to check that the object serialisation works correctly."""

35
tests/unit_tests/_primaite/_simulator/_network/test_container.py
View File

@@ -31,7 +31,6 @@ def network(example_network) -> Network:
assert len(example_network.computers) is 2
assert len(example_network.servers) is 2
example_network.set_original_state()
example_network.show()
return example_network
@@ -45,40 +44,6 @@ def test_describe_state(network):
assert len(state["links"]) is 6
def test_reset_network(network):
"""
Test that the network is properly reset.
TODO: make sure that once implemented - any installed/uninstalled services, processes, apps,
etc are also removed/reinstalled
"""
state_before = network.describe_state()
client_1: Computer = network.get_node_by_hostname("client_1")
server_1: Computer = network.get_node_by_hostname("server_1")
assert client_1.operating_state is NodeOperatingState.ON
assert server_1.operating_state is NodeOperatingState.ON
client_1.power_off()
assert client_1.operating_state is NodeOperatingState.SHUTTING_DOWN
server_1.power_off()
assert server_1.operating_state is NodeOperatingState.SHUTTING_DOWN
assert network.describe_state() != state_before
network.reset_component_for_episode(episode=1)
assert client_1.operating_state is NodeOperatingState.ON
assert server_1.operating_state is NodeOperatingState.ON
# don't worry if UUIDs change
a = filter_keys_nested_item(json.dumps(network.describe_state(), sort_keys=True, indent=2), ["uuid"])
b = filter_keys_nested_item(json.dumps(state_before, sort_keys=True, indent=2), ["uuid"])
assert a == b
def test_creating_container():
"""Check that we can create a network container"""
net = Network()

30
tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_dos_bot.py
View File

@@ -19,7 +19,6 @@ def dos_bot() -> DoSBot:
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
dos_bot.configure(target_ip_address=IPv4Address("192.168.0.1"))
dos_bot.set_original_state()
return dos_bot
@@ -28,35 +27,6 @@ def test_dos_bot_creation(dos_bot):
assert dos_bot is not None
def test_dos_bot_reset(dos_bot):
assert dos_bot.target_ip_address == IPv4Address("192.168.0.1")
assert dos_bot.target_port is Port.POSTGRES_SERVER
assert dos_bot.payload is None
assert dos_bot.repeat is False
dos_bot.configure(
target_ip_address=IPv4Address("192.168.1.1"), target_port=Port.HTTP, payload="payload", repeat=True
)
# should reset the relevant items
dos_bot.reset_component_for_episode(episode=0)
assert dos_bot.target_ip_address == IPv4Address("192.168.0.1")
assert dos_bot.target_port is Port.POSTGRES_SERVER
assert dos_bot.payload is None
assert dos_bot.repeat is False
dos_bot.configure(
target_ip_address=IPv4Address("192.168.1.1"), target_port=Port.HTTP, payload="payload", repeat=True
)
dos_bot.set_original_state()
dos_bot.reset_component_for_episode(episode=1)
# should reset to the configured value
assert dos_bot.target_ip_address == IPv4Address("192.168.1.1")
assert dos_bot.target_port is Port.HTTP
assert dos_bot.payload == "payload"
assert dos_bot.repeat is True
def test_dos_bot_cannot_run_when_node_offline(dos_bot):
dos_bot_node: Computer = dos_bot.parent
assert dos_bot_node.operating_state is NodeOperatingState.ON

12
tests/unit_tests/_primaite/_simulator/_system/test_software.py
View File

@@ -2,12 +2,14 @@ from typing import Dict
import pytest
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.software import Software, SoftwareHealthState
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import IOSoftware, SoftwareHealthState
class TestSoftware(Software):
class TestSoftware(Service):
def describe_state(self) -> Dict:
pass
@@ -15,7 +17,11 @@ class TestSoftware(Software):
@pytest.fixture(scope="function")
def software(file_system):
return TestSoftware(
name="TestSoftware", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="test_service")
name="TestSoftware",
port=Port.ARP,
file_system=file_system,
sys_log=SysLog(hostname="test_service"),
protocol=IPProtocol.TCP,
)