Merged PR 321: CAOS 0.8 observations
## Summary * Remove the usecase-specific and agent-specific observation classes, replacing with a more flexible system * Add configuration schemas to every observation class * Add router, firewall, port, and application observation * Re-shape the dict structure of observations to make it adhere to CAOS 0.8 * Change existing configs to use the new structure * make host observation separate ## Test process existing and new unit tests as well as ad hoc notebooks ## Checklist - [ ] PR is linked to a **work item** - [ ] **acceptance criteria** of linked ticket are met - [ ] performed **self-review** of the code - [ ] written **tests** for any new functionality added with this PR - [ ] updated the **documentation** if this PR changes or adds functionality - [ ] written/updated **design docs** if this PR implements new functionality - [ ] updated the **change log** - [ ] ran **pre-commit** checks for code style - [ ] attended to any **TO-DOs** left in the code Related work items: #2417
This commit is contained in:
Marek Wolan
@@ -1,6 +1,6 @@
|
||||
import pytest
|
||||
|
||||
from primaite.game.agent.observations.observations import AclObservation
|
||||
from primaite.game.agent.observations.acl_observation import ACLObservation
|
||||
from primaite.simulator.network.hardware.nodes.host.computer import Computer
|
||||
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
@@ -34,11 +34,13 @@ def test_acl_observations(simulation):
|
||||
# add router acl rule
|
||||
router.acl.add_rule(action=ACLAction.PERMIT, dst_port=Port.NTP, src_port=Port.NTP, position=1)
|
||||
|
||||
acl_obs = AclObservation(
|
||||
acl_obs = ACLObservation(
|
||||
where=["network", "nodes", router.hostname, "acl", "acl"],
|
||||
node_ip_to_id={},
|
||||
ports=["NTP", "HTTP", "POSTGRES_SERVER"],
|
||||
protocols=["TCP", "UDP", "ICMP"],
|
||||
ip_list=[],
|
||||
port_list=["NTP", "HTTP", "POSTGRES_SERVER"],
|
||||
protocol_list=["TCP", "UDP", "ICMP"],
|
||||
num_rules=10,
|
||||
wildcard_list=[],
|
||||
)
|
||||
|
||||
observation_space = acl_obs.observe(simulation.describe_state())
|
||||
@@ -46,11 +48,11 @@ def test_acl_observations(simulation):
|
||||
rule_obs = observation_space.get(1) # this is the ACL Rule added to allow NTP
|
||||
assert rule_obs.get("position") == 0 # rule was put at position 1 (0 because counting from 1 instead of 1)
|
||||
assert rule_obs.get("permission") == 1 # permit = 1 deny = 2
|
||||
assert rule_obs.get("source_node_id") == 1 # applies to all source nodes
|
||||
assert rule_obs.get("dest_node_id") == 1 # applies to all destination nodes
|
||||
assert rule_obs.get("source_port") == 2 # NTP port is mapped to value 2 (1 = ALL, so 1+1 = 2 quik mafs)
|
||||
assert rule_obs.get("dest_port") == 2 # NTP port is mapped to value 2
|
||||
assert rule_obs.get("protocol") == 1 # 1 = No Protocol
|
||||
assert rule_obs.get("source_ip_id") == 1 # applies to all source nodes
|
||||
assert rule_obs.get("dest_ip_id") == 1 # applies to all destination nodes
|
||||
assert rule_obs.get("source_port_id") == 2 # NTP port is mapped to value 2 (1 = ALL, so 1+1 = 2 quik mafs)
|
||||
assert rule_obs.get("dest_port_id") == 2 # NTP port is mapped to value 2
|
||||
assert rule_obs.get("protocol_id") == 1 # 1 = No Protocol
|
||||
|
||||
router.acl.remove_rule(1)
|
||||
|
||||
@@ -59,8 +61,8 @@ def test_acl_observations(simulation):
|
||||
rule_obs = observation_space.get(1) # this is the ACL Rule added to allow NTP
|
||||
assert rule_obs.get("position") == 0
|
||||
assert rule_obs.get("permission") == 0
|
||||
assert rule_obs.get("source_node_id") == 0
|
||||
assert rule_obs.get("dest_node_id") == 0
|
||||
assert rule_obs.get("source_port") == 0
|
||||
assert rule_obs.get("dest_port") == 0
|
||||
assert rule_obs.get("protocol") == 0
|
||||
assert rule_obs.get("source_ip_id") == 0
|
||||
assert rule_obs.get("dest_ip_id") == 0
|
||||
assert rule_obs.get("source_port_id") == 0
|
||||
assert rule_obs.get("dest_port_id") == 0
|
||||
assert rule_obs.get("protocol_id") == 0
|
||||
|
||||
@@ -23,7 +23,8 @@ def test_file_observation(simulation):
|
||||
file = pc.file_system.create_file(file_name="dog.png")
|
||||
|
||||
dog_file_obs = FileObservation(
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"]
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"],
|
||||
include_num_access=False,
|
||||
)
|
||||
|
||||
assert dog_file_obs.space["health_status"] == spaces.Discrete(6)
|
||||
@@ -49,7 +50,10 @@ def test_folder_observation(simulation):
|
||||
file = pc.file_system.create_file(file_name="dog.png", folder_name="test_folder")
|
||||
|
||||
root_folder_obs = FolderObservation(
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "test_folder"]
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "test_folder"],
|
||||
include_num_access=False,
|
||||
num_files=1,
|
||||
files=[],
|
||||
)
|
||||
|
||||
assert root_folder_obs.space["health_status"] == spaces.Discrete(6)
|
||||
@@ -68,3 +72,6 @@ def test_folder_observation(simulation):
|
||||
|
||||
observation_state = root_folder_obs.observe(simulation.describe_state())
|
||||
assert observation_state.get("health_status") == 3 # file is corrupt therefore folder is corrupted too
|
||||
|
||||
|
||||
# TODO: Add tests to check num access is correct.
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
from primaite.game.agent.observations.firewall_observation import FirewallObservation
|
||||
from primaite.simulator.network.container import Network
|
||||
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
|
||||
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
|
||||
from primaite.simulator.network.hardware.nodes.network.router import ACLAction
|
||||
from primaite.simulator.network.hardware.nodes.network.switch import Switch
|
||||
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
|
||||
|
||||
def check_default_rules(acl_obs):
|
||||
assert len(acl_obs) == 7
|
||||
assert all(acl_obs[i]["position"] == i - 1 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["permission"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_ip_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_wildcard_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_port_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_ip_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_wildcard_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_port_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["protocol_id"] == 0 for i in range(1, 8))
|
||||
|
||||
|
||||
def test_firewall_observation():
|
||||
"""Test adding/removing acl rules and enabling/disabling ports."""
|
||||
net = Network()
|
||||
firewall = Firewall(hostname="firewall", operating_state=NodeOperatingState.ON)
|
||||
firewall_observation = FirewallObservation(
|
||||
where=[],
|
||||
num_rules=7,
|
||||
ip_list=["10.0.0.1", "10.0.0.2"],
|
||||
wildcard_list=["0.0.0.255", "0.0.0.1"],
|
||||
port_list=["HTTP", "DNS"],
|
||||
protocol_list=["TCP"],
|
||||
)
|
||||
|
||||
observation = firewall_observation.observe(firewall.describe_state())
|
||||
assert "ACL" in observation
|
||||
assert "PORTS" in observation
|
||||
assert "INTERNAL" in observation["ACL"]
|
||||
assert "EXTERNAL" in observation["ACL"]
|
||||
assert "DMZ" in observation["ACL"]
|
||||
assert "INBOUND" in observation["ACL"]["INTERNAL"]
|
||||
assert "OUTBOUND" in observation["ACL"]["INTERNAL"]
|
||||
assert "INBOUND" in observation["ACL"]["EXTERNAL"]
|
||||
assert "OUTBOUND" in observation["ACL"]["EXTERNAL"]
|
||||
assert "INBOUND" in observation["ACL"]["DMZ"]
|
||||
assert "OUTBOUND" in observation["ACL"]["DMZ"]
|
||||
all_acls = (
|
||||
observation["ACL"]["INTERNAL"]["INBOUND"],
|
||||
observation["ACL"]["INTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["INBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["DMZ"]["INBOUND"],
|
||||
observation["ACL"]["DMZ"]["OUTBOUND"],
|
||||
)
|
||||
for acl_obs in all_acls:
|
||||
check_default_rules(acl_obs)
|
||||
|
||||
# add a rule to the internal inbound and check that the observation is correct
|
||||
firewall.internal_inbound_acl.add_rule(
|
||||
action=ACLAction.DENY,
|
||||
protocol=IPProtocol.TCP,
|
||||
src_ip_address="10.0.0.1",
|
||||
src_wildcard_mask="0.0.0.1",
|
||||
dst_ip_address="10.0.0.2",
|
||||
dst_wildcard_mask="0.0.0.1",
|
||||
src_port=Port.HTTP,
|
||||
dst_port=Port.HTTP,
|
||||
position=5,
|
||||
)
|
||||
|
||||
observation = firewall_observation.observe(firewall.describe_state())
|
||||
observed_rule = observation["ACL"]["INTERNAL"]["INBOUND"][5]
|
||||
assert observed_rule["position"] == 4
|
||||
assert observed_rule["permission"] == 2
|
||||
assert observed_rule["source_ip_id"] == 2
|
||||
assert observed_rule["source_wildcard_id"] == 3
|
||||
assert observed_rule["source_port_id"] == 2
|
||||
assert observed_rule["dest_ip_id"] == 3
|
||||
assert observed_rule["dest_wildcard_id"] == 3
|
||||
assert observed_rule["dest_port_id"] == 2
|
||||
assert observed_rule["protocol_id"] == 2
|
||||
|
||||
# check that none of the other acls have changed
|
||||
all_acls = (
|
||||
observation["ACL"]["INTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["INBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["DMZ"]["INBOUND"],
|
||||
observation["ACL"]["DMZ"]["OUTBOUND"],
|
||||
)
|
||||
for acl_obs in all_acls:
|
||||
check_default_rules(acl_obs)
|
||||
|
||||
# remove the rule and check that the observation is correct
|
||||
firewall.internal_inbound_acl.remove_rule(5)
|
||||
observation = firewall_observation.observe(firewall.describe_state())
|
||||
all_acls = (
|
||||
observation["ACL"]["INTERNAL"]["INBOUND"],
|
||||
observation["ACL"]["INTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["INBOUND"],
|
||||
observation["ACL"]["EXTERNAL"]["OUTBOUND"],
|
||||
observation["ACL"]["DMZ"]["INBOUND"],
|
||||
observation["ACL"]["DMZ"]["OUTBOUND"],
|
||||
)
|
||||
for acl_obs in all_acls:
|
||||
check_default_rules(acl_obs)
|
||||
|
||||
# check that there are three ports in the observation
|
||||
assert len(observation["PORTS"]) == 3
|
||||
|
||||
# check that the ports are all disabled
|
||||
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(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())
|
||||
assert observation["PORTS"][1]["operating_status"] == 1
|
||||
assert all(observation["PORTS"][i]["operating_status"] == 2 for i in range(2, 4))
|
||||
|
||||
# disable the port and check that the operating status is updated
|
||||
firewall.network_interface[1].disable()
|
||||
assert not firewall.network_interface[1].enabled
|
||||
observation = firewall_observation.observe(firewall.describe_state())
|
||||
assert all(observation["PORTS"][i]["operating_status"] == 2 for i in range(1, 4))
|
||||
@@ -1,11 +1,13 @@
|
||||
import pytest
|
||||
from gymnasium import spaces
|
||||
|
||||
from primaite.game.agent.observations.observations import LinkObservation
|
||||
from primaite.game.agent.observations.link_observation import LinkObservation
|
||||
from primaite.simulator.network.container import Network
|
||||
from primaite.simulator.network.hardware.base import Link, Node
|
||||
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.switch import Switch
|
||||
from primaite.simulator.sim_container import Simulation
|
||||
|
||||
|
||||
@@ -49,25 +51,44 @@ def simulation() -> Simulation:
|
||||
return sim
|
||||
|
||||
|
||||
def test_link_observation(simulation):
|
||||
"""Test the link observation."""
|
||||
# get a link
|
||||
link: Link = next(iter(simulation.network.links.values()))
|
||||
def test_link_observation():
|
||||
"""Check the shape and contents of the link observation."""
|
||||
net = Network()
|
||||
sim = Simulation(network=net)
|
||||
switch = Switch(hostname="switch", num_ports=5, operating_state=NodeOperatingState.ON)
|
||||
computer_1 = Computer(
|
||||
hostname="computer_1", ip_address="10.0.0.1", subnet_mask="255.255.255.0", start_up_duration=0
|
||||
)
|
||||
computer_2 = 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()
|
||||
link_1 = net.connect(switch.network_interface[1], computer_1.network_interface[1])
|
||||
link_2 = net.connect(switch.network_interface[2], computer_2.network_interface[1])
|
||||
assert link_1 is not None
|
||||
assert link_2 is not None
|
||||
|
||||
computer: Computer = simulation.network.get_node_by_hostname("computer")
|
||||
server: Server = simulation.network.get_node_by_hostname("server")
|
||||
link_1_observation = LinkObservation(where=["network", "links", link_1.uuid])
|
||||
link_2_observation = LinkObservation(where=["network", "links", link_2.uuid])
|
||||
|
||||
simulation.apply_timestep(0) # some pings when network was made - reset with apply timestep
|
||||
state = sim.describe_state()
|
||||
link_1_obs = link_1_observation.observe(state)
|
||||
link_2_obs = link_2_observation.observe(state)
|
||||
assert "PROTOCOLS" in link_1_obs
|
||||
assert "PROTOCOLS" in link_2_obs
|
||||
assert "ALL" in link_1_obs["PROTOCOLS"]
|
||||
assert "ALL" in link_2_obs["PROTOCOLS"]
|
||||
assert link_1_observation.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11)
|
||||
assert link_2_observation.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11)
|
||||
assert link_1_obs["PROTOCOLS"]["ALL"] == 0
|
||||
assert link_2_obs["PROTOCOLS"]["ALL"] == 0
|
||||
|
||||
link_obs = LinkObservation(where=["network", "links", link.uuid])
|
||||
|
||||
assert link_obs.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11) # test that the spaces are 0-10 including 0 and 10
|
||||
|
||||
observation_state = link_obs.observe(simulation.describe_state())
|
||||
assert observation_state.get("PROTOCOLS") is not None
|
||||
assert observation_state["PROTOCOLS"]["ALL"] == 0
|
||||
|
||||
computer.ping(server.network_interface.get(1).ip_address)
|
||||
|
||||
observation_state = link_obs.observe(simulation.describe_state())
|
||||
assert observation_state["PROTOCOLS"]["ALL"] == 1
|
||||
# Test that the link observation is updated when a packet is sent
|
||||
computer_1.ping("10.0.0.2")
|
||||
computer_2.ping("10.0.0.1")
|
||||
state = sim.describe_state()
|
||||
link_1_obs = link_1_observation.observe(state)
|
||||
link_2_obs = link_2_observation.observe(state)
|
||||
assert link_1_obs["PROTOCOLS"]["ALL"] > 0
|
||||
assert link_2_obs["PROTOCOLS"]["ALL"] > 0
|
||||
|
||||
@@ -5,7 +5,7 @@ import pytest
|
||||
import yaml
|
||||
from gymnasium import spaces
|
||||
|
||||
from primaite.game.agent.observations.nic_observations import NicObservation
|
||||
from primaite.game.agent.observations.nic_observations import NICObservation
|
||||
from primaite.game.game import PrimaiteGame
|
||||
from primaite.simulator.network.hardware.nodes.host.computer import Computer
|
||||
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
|
||||
@@ -40,7 +40,7 @@ def test_nic(simulation):
|
||||
|
||||
nic: NIC = pc.network_interface[1]
|
||||
|
||||
nic_obs = NicObservation(where=["network", "nodes", pc.hostname, "NICs", 1])
|
||||
nic_obs = NICObservation(where=["network", "nodes", pc.hostname, "NICs", 1], include_nmne=True)
|
||||
|
||||
assert nic_obs.space["nic_status"] == spaces.Discrete(3)
|
||||
assert nic_obs.space["NMNE"]["inbound"] == spaces.Discrete(4)
|
||||
@@ -61,17 +61,22 @@ def test_nic_categories(simulation):
|
||||
"""Test the NIC observation nmne count categories."""
|
||||
pc: Computer = simulation.network.get_node_by_hostname("client_1")
|
||||
|
||||
nic_obs = NicObservation(where=["network", "nodes", pc.hostname, "NICs", 1])
|
||||
nic_obs = NICObservation(where=["network", "nodes", pc.hostname, "NICs", 1], include_nmne=True)
|
||||
|
||||
assert nic_obs.high_nmne_threshold == 10 # default
|
||||
assert nic_obs.med_nmne_threshold == 5 # default
|
||||
assert nic_obs.low_nmne_threshold == 0 # default
|
||||
|
||||
nic_obs = NicObservation(
|
||||
|
||||
@pytest.mark.skip(reason="Feature not implemented yet")
|
||||
def test_config_nic_categories(simulation):
|
||||
pc: Computer = simulation.network.get_node_by_hostname("client_1")
|
||||
nic_obs = NICObservation(
|
||||
where=["network", "nodes", pc.hostname, "NICs", 1],
|
||||
low_nmne_threshold=3,
|
||||
med_nmne_threshold=6,
|
||||
high_nmne_threshold=9,
|
||||
include_nmne=True,
|
||||
)
|
||||
|
||||
assert nic_obs.high_nmne_threshold == 9
|
||||
@@ -80,18 +85,20 @@ def test_nic_categories(simulation):
|
||||
|
||||
with pytest.raises(Exception):
|
||||
# should throw an error
|
||||
NicObservation(
|
||||
NICObservation(
|
||||
where=["network", "nodes", pc.hostname, "NICs", 1],
|
||||
low_nmne_threshold=9,
|
||||
med_nmne_threshold=6,
|
||||
high_nmne_threshold=9,
|
||||
include_nmne=True,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
# should throw an error
|
||||
NicObservation(
|
||||
NICObservation(
|
||||
where=["network", "nodes", pc.hostname, "NICs", 1],
|
||||
low_nmne_threshold=3,
|
||||
med_nmne_threshold=9,
|
||||
high_nmne_threshold=9,
|
||||
include_nmne=True,
|
||||
)
|
||||
|
||||
@@ -4,7 +4,7 @@ from uuid import uuid4
|
||||
import pytest
|
||||
from gymnasium import spaces
|
||||
|
||||
from primaite.game.agent.observations.node_observations import NodeObservation
|
||||
from primaite.game.agent.observations.host_observations import HostObservation
|
||||
from primaite.simulator.network.hardware.nodes.host.computer import Computer
|
||||
from primaite.simulator.sim_container import Simulation
|
||||
|
||||
@@ -19,15 +19,28 @@ def simulation(example_network) -> Simulation:
|
||||
return sim
|
||||
|
||||
|
||||
def test_node_observation(simulation):
|
||||
"""Test a Node observation."""
|
||||
def test_host_observation(simulation):
|
||||
"""Test a Host observation."""
|
||||
pc: Computer = simulation.network.get_node_by_hostname("client_1")
|
||||
|
||||
node_obs = NodeObservation(where=["network", "nodes", pc.hostname])
|
||||
host_obs = HostObservation(
|
||||
where=["network", "nodes", pc.hostname],
|
||||
num_applications=0,
|
||||
num_files=1,
|
||||
num_folders=1,
|
||||
num_nics=2,
|
||||
num_services=1,
|
||||
include_num_access=False,
|
||||
include_nmne=False,
|
||||
services=[],
|
||||
applications=[],
|
||||
folders=[],
|
||||
network_interfaces=[],
|
||||
)
|
||||
|
||||
assert node_obs.space["operating_status"] == spaces.Discrete(5)
|
||||
assert host_obs.space["operating_status"] == spaces.Discrete(5)
|
||||
|
||||
observation_state = node_obs.observe(simulation.describe_state())
|
||||
observation_state = host_obs.observe(simulation.describe_state())
|
||||
assert observation_state.get("operating_status") == 1 # computer is on
|
||||
|
||||
assert observation_state.get("SERVICES") is not None
|
||||
@@ -36,11 +49,11 @@ def test_node_observation(simulation):
|
||||
|
||||
# turn off computer
|
||||
pc.power_off()
|
||||
observation_state = node_obs.observe(simulation.describe_state())
|
||||
observation_state = host_obs.observe(simulation.describe_state())
|
||||
assert observation_state.get("operating_status") == 4 # shutting down
|
||||
|
||||
for i in range(pc.shut_down_duration + 1):
|
||||
pc.apply_timestep(i)
|
||||
|
||||
observation_state = node_obs.observe(simulation.describe_state())
|
||||
observation_state = host_obs.observe(simulation.describe_state())
|
||||
assert observation_state.get("operating_status") == 2
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
from pprint import pprint
|
||||
|
||||
from primaite.game.agent.observations.acl_observation import ACLObservation
|
||||
from primaite.game.agent.observations.nic_observations import PortObservation
|
||||
from primaite.game.agent.observations.router_observation import RouterObservation
|
||||
from primaite.simulator.network.container import Network
|
||||
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
|
||||
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
|
||||
from primaite.simulator.network.hardware.nodes.network.switch import Switch
|
||||
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
from primaite.simulator.sim_container import Simulation
|
||||
|
||||
|
||||
def test_router_observation():
|
||||
"""Test adding/removing acl rules and enabling/disabling ports."""
|
||||
net = Network()
|
||||
router = Router(hostname="router", num_ports=5, operating_state=NodeOperatingState.ON)
|
||||
|
||||
ports = [PortObservation(where=["NICs", i]) for i in range(1, 6)]
|
||||
acl = ACLObservation(
|
||||
where=["acl", "acl"],
|
||||
num_rules=7,
|
||||
ip_list=["10.0.0.1", "10.0.0.2"],
|
||||
wildcard_list=["0.0.0.255", "0.0.0.1"],
|
||||
port_list=["HTTP", "DNS"],
|
||||
protocol_list=["TCP"],
|
||||
)
|
||||
router_observation = RouterObservation(where=[], ports=ports, num_ports=8, acl=acl)
|
||||
|
||||
# Observe the state using the RouterObservation instance
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
|
||||
# Check that the right number of ports and acls are in the router observation
|
||||
assert len(observed_output["PORTS"]) == 8
|
||||
assert len(observed_output["ACL"]) == 7
|
||||
|
||||
# Add an ACL rule to the router
|
||||
router.acl.add_rule(
|
||||
action=ACLAction.DENY,
|
||||
protocol=IPProtocol.TCP,
|
||||
src_ip_address="10.0.0.1",
|
||||
src_wildcard_mask="0.0.0.1",
|
||||
dst_ip_address="10.0.0.2",
|
||||
dst_wildcard_mask="0.0.0.1",
|
||||
src_port=Port.HTTP,
|
||||
dst_port=Port.HTTP,
|
||||
position=5,
|
||||
)
|
||||
# Observe the state using the RouterObservation instance
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
observed_rule = observed_output["ACL"][5]
|
||||
assert observed_rule["position"] == 4
|
||||
assert observed_rule["permission"] == 2
|
||||
assert observed_rule["source_ip_id"] == 2
|
||||
assert observed_rule["source_wildcard_id"] == 3
|
||||
assert observed_rule["source_port_id"] == 2
|
||||
assert observed_rule["dest_ip_id"] == 3
|
||||
assert observed_rule["dest_wildcard_id"] == 3
|
||||
assert observed_rule["dest_port_id"] == 2
|
||||
assert observed_rule["protocol_id"] == 2
|
||||
|
||||
# Add an ACL rule with ALL/NONE values and check that the observation is correct
|
||||
router.acl.add_rule(
|
||||
action=ACLAction.PERMIT,
|
||||
protocol=None,
|
||||
src_ip_address=None,
|
||||
src_wildcard_mask=None,
|
||||
dst_ip_address=None,
|
||||
dst_wildcard_mask=None,
|
||||
src_port=None,
|
||||
dst_port=None,
|
||||
position=2,
|
||||
)
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
observed_rule = observed_output["ACL"][2]
|
||||
assert observed_rule["position"] == 1
|
||||
assert observed_rule["permission"] == 1
|
||||
assert observed_rule["source_ip_id"] == 1
|
||||
assert observed_rule["source_wildcard_id"] == 1
|
||||
assert observed_rule["source_port_id"] == 1
|
||||
assert observed_rule["dest_ip_id"] == 1
|
||||
assert observed_rule["dest_wildcard_id"] == 1
|
||||
assert observed_rule["dest_port_id"] == 1
|
||||
assert observed_rule["protocol_id"] == 1
|
||||
|
||||
# Check that the router ports are all disabled
|
||||
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(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())
|
||||
assert observed_output["PORTS"][1]["operating_status"] == 1
|
||||
assert all(observed_output["PORTS"][i]["operating_status"] == 2 for i in range(2, 6))
|
||||
|
||||
# disable the port and check that the operating status is updated
|
||||
router.network_interface[1].disable()
|
||||
assert not router.network_interface[1].enabled
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
assert all(observed_output["PORTS"][i]["operating_status"] == 2 for i in range(1, 6))
|
||||
|
||||
# Check that ports that are out of range are shown as unused
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
assert observed_output["PORTS"][6]["operating_status"] == 0
|
||||
assert observed_output["PORTS"][7]["operating_status"] == 0
|
||||
assert observed_output["PORTS"][8]["operating_status"] == 0
|
||||
@@ -14,7 +14,13 @@ def test_file_observation():
|
||||
state = sim.describe_state()
|
||||
|
||||
dog_file_obs = FileObservation(
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"]
|
||||
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"],
|
||||
include_num_access=False,
|
||||
)
|
||||
assert dog_file_obs.observe(state) == {"health_status": 1}
|
||||
assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)})
|
||||
|
||||
|
||||
# TODO:
|
||||
# def test_file_num_access():
|
||||
# ...
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from primaite.game.agent.observations.nic_observations import NicObservation
|
||||
from primaite.game.agent.observations.nic_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
|
||||
@@ -141,9 +141,9 @@ def test_describe_state_nmne(uc2_network):
|
||||
|
||||
def test_capture_nmne_observations(uc2_network):
|
||||
"""
|
||||
Tests the NicObservation class's functionality within a simulated network environment.
|
||||
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
|
||||
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
|
||||
@@ -168,8 +168,8 @@ def test_capture_nmne_observations(uc2_network):
|
||||
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])
|
||||
db_server_nic_obs = NICObservation(where=["network", "nodes", "database_server", "NICs", 1], include_nmne=True)
|
||||
web_server_nic_obs = NICObservation(where=["network", "nodes", "web_server", "NICs", 1], include_nmne=True)
|
||||
|
||||
# Iterate through a set of test cases to simulate multiple DELETE queries
|
||||
for i in range(0, 20):
|
||||
|
||||
Reference in New Issue
Block a user