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

Merge remote-tracking branch 'origin/dev' into feature/2257-router-routes-cannot-be-represented-in-config-file

Browse Source
This commit is contained in:
Czar Echavez
2024-02-29 11:34:56 +00:00
parent 9a4587155b 312b5c35bd
commit eefc2739c8
58 changed files with 548 additions and 963 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

@@ -529,6 +529,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)

42
tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py
View File

@@ -14,7 +14,47 @@ from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
from primaite.simulator.system.services.web_server.web_server import WebServer
from tests.integration_tests.configuration_file_parsing import BASIC_CONFIG, load_config
from tests import TEST_ASSETS_ROOT
BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
"""Returns a PrimaiteGame object which loads the contents of a given yaml path."""
with open(config_path, "r") as f:
cfg = yaml.safe_load(f)
return PrimaiteGame.from_config(cfg)
def test_example_config():
"""Test that the example config can be parsed properly."""
game = load_config(example_config_path())
assert len(game.agents) == 4 # red, blue and 2 green agents
# green agent 1
assert "client_2_green_user" in game.agents
assert isinstance(game.agents["client_2_green_user"], RandomAgent)
# green agent 2
assert "client_1_green_user" in game.agents
assert isinstance(game.agents["client_1_green_user"], RandomAgent)
# red agent
assert "client_1_data_manipulation_red_bot" in game.agents
assert isinstance(game.agents["client_1_data_manipulation_red_bot"], DataManipulationAgent)
# blue agent
assert "defender" in game.agents
assert isinstance(game.agents["defender"], ProxyAgent)
network: Network = game.simulation.network
assert len(network.nodes) == 10 # 10 nodes in example network
assert len(network.routers) == 1 # 1 router in network
assert len(network.switches) == 2 # 2 switches in network
assert len(network.servers) == 5 # 5 servers in network
def test_node_software_install():

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