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

#2887 - Node class changes to address some test failures. Addressed some inconsistencies around operating_state, amended instantiation of some Nodes in test environments

Browse Source
This commit is contained in:
Charlie Crane
2025-01-27 16:35:40 +00:00
parent a7395c466e
commit 0570ab984d
46 changed files with 548 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/integration_tests/component_creation/test_action_integration.py
View File

@@ -12,12 +12,18 @@ def test_passing_actions_down(monkeypatch) -> None:
sim = Simulation()
pc1 = Computer.from_config(config={"type":"computer", "hostname":"PC-1", "ip_address":"10.10.1.1", "subnet_mask":"255.255.255.0"})
pc1 = Computer.from_config(
config={"type": "computer", "hostname": "PC-1", "ip_address": "10.10.1.1", "subnet_mask": "255.255.255.0"}
)
pc1.start_up_duration = 0
pc1.power_on()
pc2 = Computer.from_config(config={"type":"computer", "hostname":"PC-2", "ip_address":"10.10.1.2", "subnet_mask":"255.255.255.0"})
srv = Server.from_config(config={"type":"server", "hostname":"WEBSERVER", "ip_address":"10.10.1.100", "subnet_mask":"255.255.255.0"})
s1 = Switch.from_config(config={"type":"switch", "hostname":"switch1"})
pc2 = Computer.from_config(
config={"type": "computer", "hostname": "PC-2", "ip_address": "10.10.1.2", "subnet_mask": "255.255.255.0"}
)
srv = Server.from_config(
config={"type": "server", "hostname": "WEBSERVER", "ip_address": "10.10.1.100", "subnet_mask": "255.255.255.0"}
)
s1 = Switch.from_config(config={"type": "switch", "hostname": "switch1"})
for n in [pc1, pc2, srv, s1]:
sim.network.add_node(n)

2
tests/integration_tests/configuration_file_parsing/nodes/network/test_router_config.py
View File

@@ -5,8 +5,8 @@ from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import PORT_LOOKUP
from tests.integration_tests.configuration_file_parsing import DMZ_NETWORK, load_config

2
tests/integration_tests/configuration_file_parsing/nodes/test_node_config.py
View File

@@ -3,6 +3,8 @@ from primaite.config.load import data_manipulation_config_path
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from tests.integration_tests.configuration_file_parsing import BASIC_CONFIG, DMZ_NETWORK, load_config

1
tests/integration_tests/configuration_file_parsing/test_episode_scheduler.py
View File

@@ -4,6 +4,7 @@ import yaml
from primaite.session.environment import PrimaiteGymEnv
from primaite.session.ray_envs import PrimaiteRayEnv, PrimaiteRayMARLEnv
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from tests.conftest import TEST_ASSETS_ROOT
folder_path = TEST_ASSETS_ROOT / "configs" / "scenario_with_placeholders"

6
tests/integration_tests/game_layer/actions/test_node_request_permission.py
View File

@@ -35,7 +35,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN
for i in range(client_1.shut_down_duration + 1):
for i in range(client_1.config.shut_down_duration + 1):
action = ("do_nothing", {})
agent.store_action(action)
game.step()
@@ -49,7 +49,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert client_1.operating_state == NodeOperatingState.BOOTING
for i in range(client_1.start_up_duration + 1):
for i in range(client_1.config.start_up_duration + 1):
action = ("do_nothing", {})
agent.store_action(action)
game.step()
@@ -79,7 +79,7 @@ def test_node_cannot_be_shut_down_if_node_is_already_off(game_and_agent_fixture:
client_1 = game.simulation.network.get_node_by_hostname("client_1")
client_1.power_off()
for i in range(client_1.shut_down_duration + 1):
for i in range(client_1.config.shut_down_duration + 1):
action = ("do_nothing", {})
agent.store_action(action)
game.step()

4
tests/integration_tests/game_layer/observations/test_firewall_observation.py
View File

@@ -117,7 +117,9 @@ def test_firewall_observation():
assert all(observation["PORTS"][i]["operating_status"] == 2 for i in range(1, 4))
# connect a switch to the firewall and check that only the correct port is updated
switch: Switch = Switch.from_config(config={"type": "switch", "hostname":"switch", "num_ports":1, "operating_state":NodeOperatingState.ON})
switch: Switch = Switch.from_config(
config={"type": "switch", "hostname": "switch", "num_ports": 1, "operating_state": NodeOperatingState.ON}
)
link = net.connect(firewall.network_interface[1], switch.network_interface[1])
assert firewall.network_interface[1].enabled
observation = firewall_observation.observe(firewall.describe_state())

24
tests/integration_tests/game_layer/observations/test_link_observations.py
View File

@@ -56,12 +56,26 @@ def test_link_observation():
"""Check the shape and contents of the link observation."""
net = Network()
sim = Simulation(network=net)
switch: Switch = Switch.from_config(config={"type":"switch", "hostname":"switch", "num_ports":5, "operating_state":NodeOperatingState.ON})
computer_1: Computer = Computer.from_config(config={"type": "computer",
"hostname":"computer_1", "ip_address":"10.0.0.1", "subnet_mask":"255.255.255.0", "start_up_duration":0}
switch: Switch = Switch.from_config(
config={"type": "switch", "hostname": "switch", "num_ports": 5, "operating_state": NodeOperatingState.ON}
)
computer_2: Computer = Computer.from_config(config={"type":"computer",
"hostname":"computer_2", "ip_address":"10.0.0.2", "subnet_mask":"255.255.255.0", "start_up_duration":0}
computer_1: Computer = Computer.from_config(
config={
"type": "computer",
"hostname": "computer_1",
"ip_address": "10.0.0.1",
"subnet_mask": "255.255.255.0",
"start_up_duration": 0,
}
)
computer_2: Computer = Computer.from_config(
config={
"type": "computer",
"hostname": "computer_2",
"ip_address": "10.0.0.2",
"subnet_mask": "255.255.255.0",
"start_up_duration": 0,
}
)
computer_1.power_on()
computer_2.power_on()

4
tests/integration_tests/game_layer/observations/test_nic_observations.py
View File

@@ -163,7 +163,9 @@ def test_nic_monitored_traffic(simulation):
pc2: Computer = simulation.network.get_node_by_hostname("client_2")
nic_obs = NICObservation(
where=["network", "nodes", pc.config.hostname, "NICs", 1], include_nmne=False, monitored_traffic=monitored_traffic
where=["network", "nodes", pc.config.hostname, "NICs", 1],
include_nmne=False,
monitored_traffic=monitored_traffic,
)
simulation.pre_timestep(0) # apply timestep to whole sim

8
tests/integration_tests/game_layer/observations/test_router_observation.py
View File

@@ -16,7 +16,9 @@ from primaite.utils.validation.port import PORT_LOOKUP
def test_router_observation():
"""Test adding/removing acl rules and enabling/disabling ports."""
net = Network()
router = Router.from_config(config={"type": "router", "hostname":"router", "num_ports":5, "operating_state":NodeOperatingState.ON})
router = Router.from_config(
config={"type": "router", "hostname": "router", "num_ports": 5, "operating_state": NodeOperatingState.ON}
)
ports = [PortObservation(where=["NICs", i]) for i in range(1, 6)]
acl = ACLObservation(
@@ -89,7 +91,9 @@ def test_router_observation():
assert all(observed_output["PORTS"][i]["operating_status"] == 2 for i in range(1, 6))
# connect a switch to the router and check that only the correct port is updated
switch: Switch = Switch.from_config(config={"type": "switch", "hostname":"switch", "num_ports":1, "operating_state":NodeOperatingState.ON})
switch: Switch = Switch.from_config(
config={"type": "switch", "hostname": "switch", "num_ports": 1, "operating_state": NodeOperatingState.ON}
)
link = net.connect(router.network_interface[1], switch.network_interface[1])
assert router.network_interface[1].enabled
observed_output = router_observation.observe(router.describe_state())

2
tests/integration_tests/game_layer/test_action_mask.py
View File

@@ -2,8 +2,8 @@
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from primaite.simulator.system.services.service import ServiceOperatingState
from tests.conftest import TEST_ASSETS_ROOT
CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"

2
tests/integration_tests/game_layer/test_actions.py
View File

@@ -17,11 +17,11 @@ from typing import Tuple
import pytest
import yaml
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.software import SoftwareHealthState

7
tests/integration_tests/game_layer/test_observations.py
View File

@@ -8,14 +8,17 @@ from primaite.simulator.sim_container import Simulation
def test_file_observation():
sim = Simulation()
pc = Computer(hostname="beep", ip_address="123.123.123.123", subnet_mask="255.255.255.0")
pc: Computer = Computer.from_config(config={"type":"computer",
"hostname":"beep",
"ip_address":"123.123.123.123",
"subnet_mask":"255.255.255.0"})
sim.network.add_node(pc)
f = pc.file_system.create_file(file_name="dog.png")
state = sim.describe_state()
dog_file_obs = FileObservation(
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"],
where=["network", "nodes", pc.config.hostname, "file_system", "folders", "root", "files", "dog.png"],
include_num_access=False,
file_system_requires_scan=False,
)

48
tests/integration_tests/network/test_broadcast.py
View File

@@ -84,24 +84,28 @@ class BroadcastTestClient(Application, identifier="BroadcastTestClient"):
def broadcast_network() -> Network:
network = Network()
client_1_cfg = {"type": "computer",
"hostname": "client_1",
"ip_address":"192.168.1.2",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration":0}
client_1_cfg = {
"type": "computer",
"hostname": "client_1",
"ip_address": "192.168.1.2",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration": 0,
}
client_1: Computer = Computer.from_config(config=client_1_cfg)
client_1.power_on()
client_1.software_manager.install(BroadcastTestClient)
application_1 = client_1.software_manager.software["BroadcastTestClient"]
application_1.run()
client_2_cfg = {"type": "computer",
"hostname": "client_2",
"ip_address":"192.168.1.3",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration":0}
client_2_cfg = {
"type": "computer",
"hostname": "client_2",
"ip_address": "192.168.1.3",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration": 0,
}
client_2: Computer = Computer.from_config(config=client_2_cfg)
client_2.power_on()
@@ -109,14 +113,16 @@ def broadcast_network() -> Network:
application_2 = client_2.software_manager.software["BroadcastTestClient"]
application_2.run()
server_1_cfg = {"type": "server",
"hostname": "server_1",
"ip_address":"192.168.1.1",
"subnet_mask": "255.255.255.0",
"default_gateway":"192.168.1.1",
"start_up_duration": 0}
server_1_cfg = {
"type": "server",
"hostname": "server_1",
"ip_address": "192.168.1.1",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration": 0,
}
server_1 :Server = Server.from_config(config=server_1_cfg)
server_1: Server = Server.from_config(config=server_1_cfg)
server_1.power_on()
@@ -124,7 +130,9 @@ def broadcast_network() -> Network:
service: BroadcastTestService = server_1.software_manager.software["BroadcastService"]
service.start()
switch_1: Switch = Switch.from_config(config={"type": "switch", "hostname":"switch_1", "num_ports":6, "start_up_duration":0})
switch_1: Switch = Switch.from_config(
config={"type": "switch", "hostname": "switch_1", "num_ports": 6, "start_up_duration": 0}
)
switch_1.power_on()
network.connect(endpoint_a=client_1.network_interface[1], endpoint_b=switch_1.network_interface[1])

4
tests/integration_tests/network/test_firewall.py
View File

@@ -41,7 +41,9 @@ def dmz_external_internal_network() -> Network:
"""
network = Network()
firewall_node: Firewall = Firewall(hostname="firewall_1", start_up_duration=0)
firewall_node: Firewall = Firewall.from_config(
config={"type": "firewall", "hostname": "firewall_1", "start_up_duration": 0}
)
firewall_node.power_on()
# configure firewall ports
firewall_node.configure_external_port(

58
tests/integration_tests/system/red_applications/test_c2_suite_integration.py
View File

@@ -34,54 +34,64 @@ def basic_network() -> Network:
# Creating two generic nodes for the C2 Server and the C2 Beacon.
node_a_cfg = {"type": "computer",
"hostname": "node_a",
"ip_address": "192.168.0.2",
"subnet_mask": "255.255.255.252",
"default_gateway": "192.168.0.1",
"start_up_duration": 0}
node_a_cfg = {
"type": "computer",
"hostname": "node_a",
"ip_address": "192.168.0.2",
"subnet_mask": "255.255.255.252",
"default_gateway": "192.168.0.1",
"start_up_duration": 0,
}
node_a: Computer = Computer.from_config(config=node_a_cfg)
node_a.power_on()
node_a.software_manager.get_open_ports()
node_a.software_manager.install(software_class=C2Server)
node_b_cfg = {"type": "computer",
"hostname": "node_b",
"ip_address": "192.168.255.2",
"subnet_mask": "255.255.255.248",
"default_gateway": "192.168.255.1",
"start_up_duration": 0}
node_b_cfg = {
"type": "computer",
"hostname": "node_b",
"ip_address": "192.168.255.2",
"subnet_mask": "255.255.255.248",
"default_gateway": "192.168.255.1",
"start_up_duration": 0,
}
node_b: Computer = Computer.from_config(config=node_b_cfg)
node_b.power_on()
node_b.software_manager.install(software_class=C2Beacon)
# Creating a generic computer for testing remote terminal connections.
node_c_cfg = {"type": "computer",
"hostname": "node_c",
"ip_address": "192.168.255.3",
"subnet_mask": "255.255.255.248",
"default_gateway": "192.168.255.1",
"start_up_duration": 0}
node_c_cfg = {
"type": "computer",
"hostname": "node_c",
"ip_address": "192.168.255.3",
"subnet_mask": "255.255.255.248",
"default_gateway": "192.168.255.1",
"start_up_duration": 0,
}
node_c: Computer = Computer.from_config(config=node_c_cfg)
node_c.power_on()
# Creating a router to sit between node 1 and node 2.
router = Router.from_config(config={"type":"router", "hostname":"router", "num_ports":3, "start_up_duration":0})
router = Router.from_config(config={"type": "router", "hostname": "router", "num_ports": 3, "start_up_duration": 0})
# Default allow all.
router.acl.add_rule(action=ACLAction.PERMIT)
router.power_on()
# Creating switches for each client.
switch_1 = Switch.from_config(config={"type":"switch", "hostname":"switch_1", "num_ports":6, "start_up_duration":0})
switch_1 = Switch.from_config(
config={"type": "switch", "hostname": "switch_1", "num_ports": 6, "start_up_duration": 0}
)
switch_1.power_on()
# Connecting the switches to the router.
router.configure_port(port=1, ip_address="192.168.0.1", subnet_mask="255.255.255.252")
network.connect(endpoint_a=router.network_interface[1], endpoint_b=switch_1.network_interface[6])
switch_2 = Switch.from_config(config={"type":"switch", "hostname":"switch_2", "num_ports":6, "start_up_duration":0})
switch_2 = Switch.from_config(
config={"type": "switch", "hostname": "switch_2", "num_ports": 6, "start_up_duration": 0}
)
switch_2.power_on()
network.connect(endpoint_a=router.network_interface[2], endpoint_b=switch_2.network_interface[6])

2
tests/integration_tests/system/test_dns_client_server.py
View File

@@ -72,7 +72,7 @@ def test_dns_client_requests_offline_dns_server(dns_client_and_dns_server):
server.power_off()
for i in range(server.shut_down_duration + 1):
for i in range(server.config.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state == NodeOperatingState.OFF

34
tests/integration_tests/system/test_service_on_node.py
View File

@@ -13,13 +13,15 @@ from primaite.simulator.system.services.service import Service, ServiceOperating
def populated_node(
service_class,
) -> Tuple[Server, Service]:
server = Server(
hostname="server",
ip_address="192.168.0.1",
subnet_mask="255.255.255.0",
start_up_duration=0,
shut_down_duration=0,
)
server_cfg = {
"type": "server",
"hostname": "server",
"ip_address": "192.168.0.1",
"subnet_mask": "255.255.255.0",
"start_up_duration": 0,
"shut_down_duration": 0,
}
server: Server = Server.from_config(config=server_cfg)
server.power_on()
server.software_manager.install(service_class)
@@ -31,14 +33,16 @@ def populated_node(
def test_service_on_offline_node(service_class):
"""Test to check that the service cannot be interacted with when node it is on is off."""
computer: Computer = Computer(
hostname="test_computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
shut_down_duration=0,
)
computer_cfg = {
"type": "computer",
"hostname": "test_computer",
"ip_address": "192.168.1.2",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",
"start_up_duration": 0,
"shut_down_duration": 0,
}
computer: Computer = Computer.from_config(config=computer_cfg)
computer.power_on()
computer.software_manager.install(service_class)

2
tests/integration_tests/system/test_web_client_server.py
View File

@@ -94,7 +94,7 @@ def test_web_page_request_from_shut_down_server(web_client_and_web_server):
server.power_off()
for i in range(server.shut_down_duration + 1):
for i in range(server.config.shut_down_duration + 1):
server.apply_timestep(timestep=i)
# node should be off