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

#2417 update observation tests and make old tests pass

Browse Source
This commit is contained in:
Marek Wolan
2024-04-01 00:54:55 +01:00
parent 0e0df1012f
commit 0ba767d2a0
22 changed files with 767 additions and 626 deletions
Download Patch File Download Diff File Expand all files Collapse all files

251
src/primaite/config/_package_data/data_manipulation_marl.yaml
View File

@@ -40,8 +40,7 @@ agents:
0: 0.3 0: 0.3
1: 0.6 1: 0.6
2: 0.1 2: 0.1
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -90,8 +89,7 @@ agents:
0: 0.3 0: 0.3
1: 0.6 1: 0.6
2: 0.1 2: 0.1
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -140,10 +138,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -179,61 +174,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: DNSServer services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: WebServer folders:
- node_hostname: database_server - folder_name: database
folders: files:
- folder_name: database - file_name: database.db
files: - hostname: backup_server
- file_name: database.db - hostname: security_suite
- node_hostname: backup_server - hostname: client_1
- node_hostname: security_suite - hostname: client_2
- node_hostname: client_1 num_services: 1
- node_hostname: client_2 num_applications: 0
links: num_folders: 1
- link_ref: router_1___switch_1 num_files: 1
- link_ref: router_1___switch_2 num_nics: 2
- link_ref: switch_1___domain_controller include_num_access: false
- link_ref: switch_1___web_server include_nmne: true
- link_ref: switch_1___database_server routers:
- link_ref: switch_1___backup_server - hostname: router_1
- link_ref: switch_1___security_suite num_ports: 0
- link_ref: switch_2___client_1 ip_list:
- link_ref: switch_2___client_2 - 192.168.1.10
- link_ref: switch_2___security_suite - 192.168.1.12
acl: - 192.168.1.14
options: - 192.168.1.16
max_acl_rules: 10 - 192.168.1.110
router_hostname: router_1 - 192.168.10.21
ip_address_order: - 192.168.10.22
- node_hostname: domain_controller - 192.168.10.110
nic_num: 1 wildcard_list:
- node_hostname: web_server - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: database_server - 80
nic_num: 1 - 5432
- node_hostname: backup_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: security_suite - TCP
nic_num: 1 - UDP
- node_hostname: client_1 num_rules: 10
nic_num: 1
- node_hostname: client_2 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: security_suite options:
nic_num: 2 link_references:
ics: null - router_1___switch_1
- router_1___switch_2
- switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:
@@ -730,61 +737,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: DNSServer services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: WebServer folders:
- node_hostname: database_server - folder_name: database
folders: files:
- folder_name: database - file_name: database.db
files: - hostname: backup_server
- file_name: database.db - hostname: security_suite
- node_hostname: backup_server - hostname: client_1
- node_hostname: security_suite - hostname: client_2
- node_hostname: client_1 num_services: 1
- node_hostname: client_2 num_applications: 0
links: num_folders: 1
- link_ref: router_1___switch_1 num_files: 1
- link_ref: router_1___switch_2 num_nics: 2
- link_ref: switch_1___domain_controller include_num_access: false
- link_ref: switch_1___web_server include_nmne: true
- link_ref: switch_1___database_server routers:
- link_ref: switch_1___backup_server - hostname: router_1
- link_ref: switch_1___security_suite num_ports: 0
- link_ref: switch_2___client_1 ip_list:
- link_ref: switch_2___client_2 - 192.168.1.10
- link_ref: switch_2___security_suite - 192.168.1.12
acl: - 192.168.1.14
options: - 192.168.1.16
max_acl_rules: 10 - 192.168.1.110
router_hostname: router_1 - 192.168.10.21
ip_address_order: - 192.168.10.22
- node_hostname: domain_controller - 192.168.10.110
nic_num: 1 wildcard_list:
- node_hostname: web_server - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: database_server - 80
nic_num: 1 - 5432
- node_hostname: backup_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: security_suite - TCP
nic_num: 1 - UDP
- node_hostname: client_1 num_rules: 10
nic_num: 1
- node_hostname: client_2 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: security_suite options:
nic_num: 2 link_references:
ics: null - router_1___switch_1
- router_1___switch_2
- switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

22
src/primaite/game/agent/observations/acl_observation.py
View File

@@ -59,10 +59,10 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
""" """
self.where = where self.where = where
self.num_rules: int = num_rules self.num_rules: int = num_rules
self.ip_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(ip_list)} self.ip_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(ip_list)}
self.wildcard_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(wildcard_list)} self.wildcard_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(wildcard_list)}
self.port_to_id: Dict[int, int] = {i + 2: p for i, p in enumerate(port_list)} self.port_to_id: Dict[int, int] = {p: i + 2 for i, p in enumerate(port_list)}
self.protocol_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(protocol_list)} self.protocol_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(protocol_list)}
self.default_observation: Dict = { self.default_observation: Dict = {
i i
+ 1: { + 1: {
@@ -110,16 +110,16 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
} }
else: else:
src_ip = rule_state["src_ip_address"] src_ip = rule_state["src_ip_address"]
src_node_id = self.ip_to_id.get(src_ip, 1) src_node_id = 1 if src_ip is None else self.ip_to_id[src_ip]
dst_ip = rule_state["dst_ip_address"] dst_ip = rule_state["dst_ip_address"]
dst_node_ip = self.ip_to_id.get(dst_ip, 1) dst_node_id = 1 if dst_ip is None else self.ip_to_id[dst_ip]
src_wildcard = rule_state["source_wildcard_id"] src_wildcard = rule_state["src_wildcard_mask"]
src_wildcard_id = self.wildcard_to_id.get(src_wildcard, 1) src_wildcard_id = self.wildcard_to_id.get(src_wildcard, 1)
dst_wildcard = rule_state["dest_wildcard_id"] dst_wildcard = rule_state["dst_wildcard_mask"]
dst_wildcard_id = self.wildcard_to_id.get(dst_wildcard, 1) dst_wildcard_id = self.wildcard_to_id.get(dst_wildcard, 1)
src_port = rule_state["source_port_id"] src_port = rule_state["src_port"]
src_port_id = self.port_to_id.get(src_port, 1) src_port_id = self.port_to_id.get(src_port, 1)
dst_port = rule_state["dest_port_id"] dst_port = rule_state["dst_port"]
dst_port_id = self.port_to_id.get(dst_port, 1) dst_port_id = self.port_to_id.get(dst_port, 1)
protocol = rule_state["protocol"] protocol = rule_state["protocol"]
protocol_id = self.protocol_to_id.get(protocol, 1) protocol_id = self.protocol_to_id.get(protocol, 1)
@@ -129,7 +129,7 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
"source_ip_id": src_node_id, "source_ip_id": src_node_id,
"source_wildcard_id": src_wildcard_id, "source_wildcard_id": src_wildcard_id,
"source_port_id": src_port_id, "source_port_id": src_port_id,
"dest_ip_id": dst_node_ip, "dest_ip_id": dst_node_id,
"dest_wildcard_id": dst_wildcard_id, "dest_wildcard_id": dst_wildcard_id,
"dest_port_id": dst_port_id, "dest_port_id": dst_port_id,
"protocol_id": protocol_id, "protocol_id": protocol_id,

16
src/primaite/game/agent/observations/file_system_observations.py
View File

@@ -133,8 +133,9 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
self.default_observation = { self.default_observation = {
"health_status": 0, "health_status": 0,
"FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)},
} }
if self.files:
self.default_observation["FILES"] = {i + 1: f.default_observation for i, f in enumerate(self.files)}
def observe(self, state: Dict) -> ObsType: def observe(self, state: Dict) -> ObsType:
""" """
@@ -154,7 +155,8 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
obs = {} obs = {}
obs["health_status"] = health_status obs["health_status"] = health_status
obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)} if self.files:
obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)}
return obs return obs
@@ -166,12 +168,10 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
:return: Gymnasium space representing the observation space for folder status. :return: Gymnasium space representing the observation space for folder status.
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict( shape = {"health_status": spaces.Discrete(6)}
{ if self.files:
"health_status": spaces.Discrete(6), shape["FILES"] = spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)})
"FILES": spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)}), return spaces.Dict(shape)
}
)
@classmethod @classmethod
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> FolderObservation: def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> FolderObservation:

52
src/primaite/game/agent/observations/host_observations.py
View File

@@ -123,21 +123,27 @@ class HostObservation(AbstractObservation, identifier="HOST"):
msg = f"Too many folders in Node observation space for node. Truncating folder {truncated_folder.where}" msg = f"Too many folders in Node observation space for node. Truncating folder {truncated_folder.where}"
_LOGGER.warning(msg) _LOGGER.warning(msg)
self.network_interfaces: List[NICObservation] = network_interfaces self.nics: List[NICObservation] = network_interfaces
while len(self.network_interfaces) < num_nics: while len(self.nics) < num_nics:
self.network_interfaces.append(NICObservation(where=None, include_nmne=include_nmne)) self.nics.append(NICObservation(where=None, include_nmne=include_nmne))
while len(self.network_interfaces) > num_nics: while len(self.nics) > num_nics:
truncated_nic = self.network_interfaces.pop() truncated_nic = self.nics.pop()
msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_nic.where}" msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_nic.where}"
_LOGGER.warning(msg) _LOGGER.warning(msg)
self.default_observation: ObsType = { self.default_observation: ObsType = {
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
"APPLICATIONS": {i + 1: a.default_observation for i, a in enumerate(self.applications)},
"FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)},
"NICS": {i + 1: n.default_observation for i, n in enumerate(self.network_interfaces)},
"operating_status": 0, "operating_status": 0,
} }
if self.services:
self.default_observation["SERVICES"] = {i + 1: s.default_observation for i, s in enumerate(self.services)}
if self.applications:
self.default_observation["APPLICATIONS"] = {
i + 1: a.default_observation for i, a in enumerate(self.applications)
}
if self.folders:
self.default_observation["FOLDERS"] = {i + 1: f.default_observation for i, f in enumerate(self.folders)}
if self.nics:
self.default_observation["NICS"] = {i + 1: n.default_observation for i, n in enumerate(self.nics)}
if self.include_num_access: if self.include_num_access:
self.default_observation["num_file_creations"] = 0 self.default_observation["num_file_creations"] = 0
self.default_observation["num_file_deletions"] = 0 self.default_observation["num_file_deletions"] = 0
@@ -156,13 +162,15 @@ class HostObservation(AbstractObservation, identifier="HOST"):
return self.default_observation return self.default_observation
obs = {} obs = {}
obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)}
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
obs["operating_status"] = node_state["operating_state"] obs["operating_status"] = node_state["operating_state"]
obs["NICS"] = { if self.services:
i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces) obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
} if self.applications:
obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)}
if self.folders:
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
if self.nics:
obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)}
if self.include_num_access: if self.include_num_access:
obs["num_file_creations"] = node_state["file_system"]["num_file_creations"] obs["num_file_creations"] = node_state["file_system"]["num_file_creations"]
obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"] obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"]
@@ -177,14 +185,16 @@ class HostObservation(AbstractObservation, identifier="HOST"):
:rtype: spaces.Space :rtype: spaces.Space
""" """
shape = { shape = {
"SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
"APPLICATIONS": spaces.Dict({i + 1: app.space for i, app in enumerate(self.applications)}),
"FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
"operating_status": spaces.Discrete(5), "operating_status": spaces.Discrete(5),
"NICS": spaces.Dict(
{i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}
),
} }
if self.services:
shape["SERVICES"] = spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)})
if self.applications:
shape["APPLICATIONS"] = spaces.Dict({i + 1: app.space for i, app in enumerate(self.applications)})
if self.folders:
shape["FOLDERS"] = spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)})
if self.nics:
shape["NICS"] = spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)})
if self.include_num_access: if self.include_num_access:
shape["num_file_creations"] = spaces.Discrete(4) shape["num_file_creations"] = spaces.Discrete(4)
shape["num_file_deletions"] = spaces.Discrete(4) shape["num_file_deletions"] = spaces.Discrete(4)

36
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -23,7 +23,11 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
include_nmne: Optional[bool] = None include_nmne: Optional[bool] = None
"""Whether to include number of malicious network events (NMNE) in the observation.""" """Whether to include number of malicious network events (NMNE) in the observation."""
def __init__(self, where: WhereType, include_nmne: bool) -> None: def __init__(
self,
where: WhereType,
include_nmne: bool,
) -> None:
""" """
Initialise a network interface observation instance. Initialise a network interface observation instance.
@@ -40,6 +44,36 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
self.default_observation: ObsType = {"nic_status": 0} self.default_observation: ObsType = {"nic_status": 0}
if self.include_nmne: if self.include_nmne:
self.default_observation.update({"NMNE": {"inbound": 0, "outbound": 0}}) self.default_observation.update({"NMNE": {"inbound": 0, "outbound": 0}})
self.nmne_inbound_last_step: int = 0
self.nmne_outbound_last_step: int = 0
# TODO: allow these to be configured in yaml
self.high_nmne_threshold = 10
self.med_nmne_threshold = 5
self.low_nmne_threshold = 0
def _categorise_mne_count(self, nmne_count: int) -> int:
"""
Categorise the number of Malicious Network Events (NMNEs) into discrete bins.
This helps in classifying the severity or volume of MNEs into manageable levels for the agent.
Bins are defined as follows:
- 0: No MNEs detected (0 events).
- 1: Low number of MNEs (default 1-5 events).
- 2: Moderate number of MNEs (default 6-10 events).
- 3: High number of MNEs (default more than 10 events).
:param nmne_count: Number of MNEs detected.
:return: Bin number corresponding to the number of MNEs. Returns 0, 1, 2, or 3 based on the detected MNE count.
"""
if nmne_count > self.high_nmne_threshold:
return 3
elif nmne_count > self.med_nmne_threshold:
return 2
elif nmne_count > self.low_nmne_threshold:
return 1
return 0
def observe(self, state: Dict) -> ObsType: def observe(self, state: Dict) -> ObsType:
""" """

13
src/primaite/game/agent/observations/router_observation.py
View File

@@ -74,9 +74,10 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
_LOGGER.warning(msg) _LOGGER.warning(msg)
self.default_observation = { self.default_observation = {
"PORTS": {i + 1: p.default_observation for i, p in enumerate(self.ports)},
"ACL": self.acl.default_observation, "ACL": self.acl.default_observation,
} }
if self.ports:
self.default_observation["PORTS"] = {i + 1: p.default_observation for i, p in enumerate(self.ports)}
def observe(self, state: Dict) -> ObsType: def observe(self, state: Dict) -> ObsType:
""" """
@@ -92,8 +93,9 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
return self.default_observation return self.default_observation
obs = {} obs = {}
obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
obs["ACL"] = self.acl.observe(state) obs["ACL"] = self.acl.observe(state)
if self.ports:
obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
return obs return obs
@property @property
@@ -104,9 +106,10 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
:return: Gymnasium space representing the observation space for router status. :return: Gymnasium space representing the observation space for router status.
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict( shape = {"ACL": self.acl.space}
{"PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}), "ACL": self.acl.space} if self.ports:
) shape["PORTS"] = spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)})
return spaces.Dict(shape)
@classmethod @classmethod
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> RouterObservation: def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> RouterObservation:

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

@@ -22,8 +22,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -50,10 +49,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -86,63 +82,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
- node_hostname: security_suite num_applications: 0
- node_hostname: client_1 num_folders: 1
- node_hostname: client_2 num_files: 1
links: num_nics: 2
- link_ref: router_1___switch_1 include_num_access: false
- link_ref: router_1___switch_2 include_nmne: true
- link_ref: switch_1___domain_controller routers:
- link_ref: switch_1___web_server - hostname: router_1
- link_ref: switch_1___database_server num_ports: 0
- link_ref: switch_1___backup_server ip_list:
- link_ref: switch_1___security_suite - 192.168.1.10
- link_ref: switch_2___client_1 - 192.168.1.12
- link_ref: switch_2___client_2 - 192.168.1.14
- link_ref: switch_2___security_suite - 192.168.1.16
acl: - 192.168.1.110
options: - 192.168.10.21
max_acl_rules: 10 - 192.168.10.22
router_hostname: router_1 - 192.168.10.110
ip_address_order: wildcard_list:
- node_hostname: domain_controller - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: web_server - 80
nic_num: 1 - 5432
- node_hostname: database_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: backup_server - TCP
nic_num: 1 - UDP
- node_hostname: security_suite num_rules: 10
nic_num: 1
- node_hostname: client_1 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_2 options:
nic_num: 1 link_references:
- node_hostname: security_suite - router_1___switch_1
nic_num: 2 - router_1___switch_2
ics: null - switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

3
tests/assets/configs/basic_firewall.yaml
View File

@@ -41,8 +41,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING

3
tests/assets/configs/basic_switched_network.yaml
View File

@@ -41,8 +41,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING

3
tests/assets/configs/dmz_network.yaml
View File

@@ -66,8 +66,7 @@ agents:
- ref: client_1_green_user - ref: client_1_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING

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

@@ -26,8 +26,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -55,10 +54,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -90,63 +86,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
- node_hostname: security_suite num_applications: 0
- node_hostname: client_1 num_folders: 1
- node_hostname: client_2 num_files: 1
links: num_nics: 2
- link_ref: router_1___switch_1 include_num_access: false
- link_ref: router_1___switch_2 include_nmne: true
- link_ref: switch_1___domain_controller routers:
- link_ref: switch_1___web_server - hostname: router_1
- link_ref: switch_1___database_server num_ports: 0
- link_ref: switch_1___backup_server ip_list:
- link_ref: switch_1___security_suite - 192.168.1.10
- link_ref: switch_2___client_1 - 192.168.1.12
- link_ref: switch_2___client_2 - 192.168.1.14
- link_ref: switch_2___security_suite - 192.168.1.16
acl: - 192.168.1.110
options: - 192.168.10.21
max_acl_rules: 10 - 192.168.10.22
router_hostname: router_1 - 192.168.10.110
ip_address_order: wildcard_list:
- node_hostname: domain_controller - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: web_server - 80
nic_num: 1 - 5432
- node_hostname: database_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: backup_server - TCP
nic_num: 1 - UDP
- node_hostname: security_suite num_rules: 10
nic_num: 1
- node_hostname: client_1 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_2 options:
nic_num: 1 link_references:
- node_hostname: security_suite - router_1___switch_1
nic_num: 2 - router_1___switch_2
ics: null - switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

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

@@ -32,8 +32,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -61,10 +60,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -97,63 +93,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
- node_hostname: security_suite num_applications: 0
- node_hostname: client_1 num_folders: 1
- node_hostname: client_2 num_files: 1
links: num_nics: 2
- link_ref: router_1___switch_1 include_num_access: false
- link_ref: router_1___switch_2 include_nmne: true
- link_ref: switch_1___domain_controller routers:
- link_ref: switch_1___web_server - hostname: router_1
- link_ref: switch_1___database_server num_ports: 0
- link_ref: switch_1___backup_server ip_list:
- link_ref: switch_1___security_suite - 192.168.1.10
- link_ref: switch_2___client_1 - 192.168.1.12
- link_ref: switch_2___client_2 - 192.168.1.14
- link_ref: switch_2___security_suite - 192.168.1.16
acl: - 192.168.1.110
options: - 192.168.10.21
max_acl_rules: 10 - 192.168.10.22
router_hostname: router_1 - 192.168.10.110
ip_address_order: wildcard_list:
- node_hostname: domain_controller - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: web_server - 80
nic_num: 1 - 5432
- node_hostname: database_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: backup_server - TCP
nic_num: 1 - UDP
- node_hostname: security_suite num_rules: 10
nic_num: 1
- node_hostname: client_1 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_2 options:
nic_num: 1 link_references:
- node_hostname: security_suite - router_1___switch_1
nic_num: 2 - router_1___switch_2
ics: null - switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:
@@ -541,63 +547,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
- node_hostname: security_suite num_applications: 0
- node_hostname: client_1 num_folders: 1
- node_hostname: client_2 num_files: 1
links: num_nics: 2
- link_ref: router_1___switch_1 include_num_access: false
- link_ref: router_1___switch_2 include_nmne: true
- link_ref: switch_1___domain_controller routers:
- link_ref: switch_1___web_server - hostname: router_1
- link_ref: switch_1___database_server num_ports: 0
- link_ref: switch_1___backup_server ip_list:
- link_ref: switch_1___security_suite - 192.168.1.10
- link_ref: switch_2___client_1 - 192.168.1.12
- link_ref: switch_2___client_2 - 192.168.1.14
- link_ref: switch_2___security_suite - 192.168.1.16
acl: - 192.168.1.110
options: - 192.168.10.21
max_acl_rules: 10 - 192.168.10.22
router_hostname: router_1 - 192.168.10.110
ip_address_order: wildcard_list:
- node_hostname: domain_controller - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: web_server - 80
nic_num: 1 - 5432
- node_hostname: database_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: backup_server - TCP
nic_num: 1 - UDP
- node_hostname: security_suite num_rules: 10
nic_num: 1
- node_hostname: client_1 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_2 options:
nic_num: 1 link_references:
- node_hostname: security_suite - router_1___switch_1
nic_num: 2 - router_1___switch_2
ics: null - switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

131
tests/assets/configs/shared_rewards.yaml
View File

@@ -41,8 +41,7 @@ agents:
0: 0.3 0: 0.3
1: 0.6 1: 0.6
2: 0.1 2: 0.1
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -91,8 +90,7 @@ agents:
0: 0.3 0: 0.3
1: 0.6 1: 0.6
2: 0.1 2: 0.1
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -141,10 +139,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -177,61 +172,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: DNSServer services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: WebServer folders:
- node_hostname: database_server - folder_name: database
folders: files:
- folder_name: database - file_name: database.db
files: - hostname: backup_server
- file_name: database.db - hostname: security_suite
- node_hostname: backup_server - hostname: client_1
- node_hostname: security_suite - hostname: client_2
- node_hostname: client_1 num_services: 1
- node_hostname: client_2 num_applications: 0
links: num_folders: 1
- link_ref: router_1___switch_1 num_files: 1
- link_ref: router_1___switch_2 num_nics: 2
- link_ref: switch_1___domain_controller include_num_access: false
- link_ref: switch_1___web_server include_nmne: true
- link_ref: switch_1___database_server routers:
- link_ref: switch_1___backup_server - hostname: router_1
- link_ref: switch_1___security_suite num_ports: 0
- link_ref: switch_2___client_1 ip_list:
- link_ref: switch_2___client_2 - 192.168.1.10
- link_ref: switch_2___security_suite - 192.168.1.12
acl: - 192.168.1.14
options: - 192.168.1.16
max_acl_rules: 10 - 192.168.1.110
router_hostname: router_1 - 192.168.10.21
ip_address_order: - 192.168.10.22
- node_hostname: domain_controller - 192.168.10.110
nic_num: 1 wildcard_list:
- node_hostname: web_server - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: database_server - 80
nic_num: 1 - 5432
- node_hostname: backup_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: security_suite - TCP
nic_num: 1 - UDP
- node_hostname: client_1 num_rules: 10
nic_num: 1
- node_hostname: client_2 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: security_suite options:
nic_num: 2 link_references:
ics: null - router_1___switch_1
- router_1___switch_2
- switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

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

@@ -33,8 +33,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -62,10 +61,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -98,65 +94,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
# services: num_applications: 0
# - service_name: backup_service num_folders: 1
- node_hostname: security_suite num_files: 1
- node_hostname: client_1 num_nics: 2
- node_hostname: client_2 include_num_access: false
links: include_nmne: true
- link_ref: router_1___switch_1 routers:
- link_ref: router_1___switch_2 - hostname: router_1
- link_ref: switch_1___domain_controller num_ports: 0
- link_ref: switch_1___web_server ip_list:
- link_ref: switch_1___database_server - 192.168.1.10
- link_ref: switch_1___backup_server - 192.168.1.12
- link_ref: switch_1___security_suite - 192.168.1.14
- link_ref: switch_2___client_1 - 192.168.1.16
- link_ref: switch_2___client_2 - 192.168.1.110
- link_ref: switch_2___security_suite - 192.168.10.21
acl: - 192.168.10.22
options: - 192.168.10.110
max_acl_rules: 10 wildcard_list:
router_hostname: router_1 - 0.0.0.1
ip_address_order: port_list:
- node_hostname: domain_controller - 80
nic_num: 1 - 5432
- node_hostname: web_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: database_server - TCP
nic_num: 1 - UDP
- node_hostname: backup_server num_rules: 10
nic_num: 1
- node_hostname: security_suite - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_1 options:
nic_num: 1 link_references:
- node_hostname: client_2 - router_1___switch_1
nic_num: 1 - router_1___switch_2
- node_hostname: security_suite - switch_1___domain_controller
nic_num: 2 - switch_1___web_server
ics: null - switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

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

@@ -26,8 +26,7 @@ agents:
- ref: client_2_green_user - ref: client_2_green_user
team: GREEN team: GREEN
type: ProbabilisticAgent type: ProbabilisticAgent
observation_space: observation_space: null
type: UC2GreenObservation
action_space: action_space:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
@@ -62,10 +61,7 @@ agents:
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
observation_space: observation_space: null
type: UC2RedObservation
options:
nodes: {}
action_space: action_space:
action_list: action_list:
@@ -98,63 +94,73 @@ agents:
type: ProxyAgent type: ProxyAgent
observation_space: observation_space:
type: UC2BlueObservation type: CUSTOM
options: options:
num_services_per_node: 1 components:
num_folders_per_node: 1 - type: NODES
num_files_per_folder: 1 label: NODES
num_nics_per_node: 2 options:
nodes: hosts:
- node_hostname: domain_controller - hostname: domain_controller
services: - hostname: web_server
- service_name: domain_controller_dns_server services:
- node_hostname: web_server - service_name: WebServer
services: - hostname: database_server
- service_name: web_server_database_client folders:
- node_hostname: database_server - folder_name: database
services: files:
- service_name: database_service - file_name: database.db
folders: - hostname: backup_server
- folder_name: database - hostname: security_suite
files: - hostname: client_1
- file_name: database.db - hostname: client_2
- node_hostname: backup_server num_services: 1
- node_hostname: security_suite num_applications: 0
- node_hostname: client_1 num_folders: 1
- node_hostname: client_2 num_files: 1
links: num_nics: 2
- link_ref: router_1___switch_1 include_num_access: false
- link_ref: router_1___switch_2 include_nmne: true
- link_ref: switch_1___domain_controller routers:
- link_ref: switch_1___web_server - hostname: router_1
- link_ref: switch_1___database_server num_ports: 0
- link_ref: switch_1___backup_server ip_list:
- link_ref: switch_1___security_suite - 192.168.1.10
- link_ref: switch_2___client_1 - 192.168.1.12
- link_ref: switch_2___client_2 - 192.168.1.14
- link_ref: switch_2___security_suite - 192.168.1.16
acl: - 192.168.1.110
options: - 192.168.10.21
max_acl_rules: 10 - 192.168.10.22
router_hostname: router_1 - 192.168.10.110
ip_address_order: wildcard_list:
- node_hostname: domain_controller - 0.0.0.1
nic_num: 1 port_list:
- node_hostname: web_server - 80
nic_num: 1 - 5432
- node_hostname: database_server protocol_list:
nic_num: 1 - ICMP
- node_hostname: backup_server - TCP
nic_num: 1 - UDP
- node_hostname: security_suite num_rules: 10
nic_num: 1
- node_hostname: client_1 - type: LINKS
nic_num: 1 label: LINKS
- node_hostname: client_2 options:
nic_num: 1 link_references:
- node_hostname: security_suite - router_1___switch_1
nic_num: 2 - router_1___switch_2
ics: null - switch_1___domain_controller
- switch_1___web_server
- switch_1___database_server
- switch_1___backup_server
- switch_1___security_suite
- switch_2___client_1
- switch_2___client_2
- switch_2___security_suite
- type: "NONE"
label: ICS
options: {}
action_space: action_space:
action_list: action_list:

8
tests/e2e_integration_tests/test_primaite_session.py
View File

@@ -11,8 +11,9 @@ MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml" MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
# @pytest.mark.skip(reason="no way of currently testing this") @pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
class TestPrimaiteSession: class TestPrimaiteSession:
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
@pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True) @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
def test_creating_session(self, temp_primaite_session): def test_creating_session(self, temp_primaite_session):
"""Check that creating a session from config works.""" """Check that creating a session from config works."""
@@ -51,6 +52,7 @@ class TestPrimaiteSession:
assert checkpoint_2.exists() assert checkpoint_2.exists()
assert not checkpoint_3.exists() assert not checkpoint_3.exists()
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
@pytest.mark.parametrize("temp_primaite_session", [[TRAINING_ONLY_PATH]], indirect=True) @pytest.mark.parametrize("temp_primaite_session", [[TRAINING_ONLY_PATH]], indirect=True)
def test_training_only_session(self, temp_primaite_session): def test_training_only_session(self, temp_primaite_session):
"""Check that you can run a training-only session.""" """Check that you can run a training-only session."""
@@ -59,6 +61,7 @@ class TestPrimaiteSession:
session.start_session() session.start_session()
# TODO: include checks that the model was trained, e.g. that the loss changed and checkpoints were saved? # TODO: include checks that the model was trained, e.g. that the loss changed and checkpoints were saved?
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
@pytest.mark.parametrize("temp_primaite_session", [[EVAL_ONLY_PATH]], indirect=True) @pytest.mark.parametrize("temp_primaite_session", [[EVAL_ONLY_PATH]], indirect=True)
def test_eval_only_session(self, temp_primaite_session): def test_eval_only_session(self, temp_primaite_session):
"""Check that you can load a model and run an eval-only session.""" """Check that you can load a model and run an eval-only session."""
@@ -67,6 +70,7 @@ class TestPrimaiteSession:
session.start_session() session.start_session()
# TODO: include checks that the model was loaded and that the eval-only session ran # TODO: include checks that the model was loaded and that the eval-only session ran
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
@pytest.mark.skip(reason="Slow, reenable later") @pytest.mark.skip(reason="Slow, reenable later")
@pytest.mark.parametrize("temp_primaite_session", [[MULTI_AGENT_PATH]], indirect=True) @pytest.mark.parametrize("temp_primaite_session", [[MULTI_AGENT_PATH]], indirect=True)
def test_multi_agent_session(self, temp_primaite_session): def test_multi_agent_session(self, temp_primaite_session):
@@ -74,10 +78,12 @@ class TestPrimaiteSession:
with temp_primaite_session as session: with temp_primaite_session as session:
session.start_session() session.start_session()
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
def test_error_thrown_on_bad_configuration(self): def test_error_thrown_on_bad_configuration(self):
with pytest.raises(pydantic.ValidationError): with pytest.raises(pydantic.ValidationError):
session = TempPrimaiteSession.from_config(MISCONFIGURED_PATH) session = TempPrimaiteSession.from_config(MISCONFIGURED_PATH)
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
@pytest.mark.skip( @pytest.mark.skip(
reason="Currently software cannot be dynamically created/destroyed during simulation. Therefore, " reason="Currently software cannot be dynamically created/destroyed during simulation. Therefore, "
"reset doesn't implement software restore." "reset doesn't implement software restore."

28
tests/integration_tests/game_layer/observations/test_acl_observations.py
View File

@@ -36,9 +36,11 @@ def test_acl_observations(simulation):
acl_obs = ACLObservation( acl_obs = ACLObservation(
where=["network", "nodes", router.hostname, "acl", "acl"], where=["network", "nodes", router.hostname, "acl", "acl"],
node_ip_to_id={}, ip_list=[],
ports=["NTP", "HTTP", "POSTGRES_SERVER"], port_list=["NTP", "HTTP", "POSTGRES_SERVER"],
protocols=["TCP", "UDP", "ICMP"], protocol_list=["TCP", "UDP", "ICMP"],
num_rules=10,
wildcard_list=[],
) )
observation_space = acl_obs.observe(simulation.describe_state()) 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 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("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("permission") == 1 # permit = 1 deny = 2
assert rule_obs.get("source_node_id") == 1 # applies to all source nodes assert rule_obs.get("source_ip_id") == 1 # applies to all source nodes
assert rule_obs.get("dest_node_id") == 1 # applies to all destination nodes assert rule_obs.get("dest_ip_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("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") == 2 # NTP port is mapped to value 2 assert rule_obs.get("dest_port_id") == 2 # NTP port is mapped to value 2
assert rule_obs.get("protocol") == 1 # 1 = No Protocol assert rule_obs.get("protocol_id") == 1 # 1 = No Protocol
router.acl.remove_rule(1) 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 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("position") == 0
assert rule_obs.get("permission") == 0 assert rule_obs.get("permission") == 0
assert rule_obs.get("source_node_id") == 0 assert rule_obs.get("source_ip_id") == 0
assert rule_obs.get("dest_node_id") == 0 assert rule_obs.get("dest_ip_id") == 0
assert rule_obs.get("source_port") == 0 assert rule_obs.get("source_port_id") == 0
assert rule_obs.get("dest_port") == 0 assert rule_obs.get("dest_port_id") == 0
assert rule_obs.get("protocol") == 0 assert rule_obs.get("protocol_id") == 0

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

@@ -23,7 +23,8 @@ def test_file_observation(simulation):
file = pc.file_system.create_file(file_name="dog.png") file = pc.file_system.create_file(file_name="dog.png")
dog_file_obs = FileObservation( 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) 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") file = pc.file_system.create_file(file_name="dog.png", folder_name="test_folder")
root_folder_obs = FolderObservation( 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) assert root_folder_obs.space["health_status"] == spaces.Discrete(6)

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

@@ -40,7 +40,7 @@ def test_nic(simulation):
nic: NIC = pc.network_interface[1] 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["nic_status"] == spaces.Discrete(3)
assert nic_obs.space["NMNE"]["inbound"] == spaces.Discrete(4) 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.""" """Test the NIC observation nmne count categories."""
pc: Computer = simulation.network.get_node_by_hostname("client_1") 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.high_nmne_threshold == 10 # default
assert nic_obs.med_nmne_threshold == 5 # default assert nic_obs.med_nmne_threshold == 5 # default
assert nic_obs.low_nmne_threshold == 0 # default assert nic_obs.low_nmne_threshold == 0 # default
@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( nic_obs = NICObservation(
where=["network", "nodes", pc.hostname, "NICs", 1], where=["network", "nodes", pc.hostname, "NICs", 1],
low_nmne_threshold=3, low_nmne_threshold=3,
med_nmne_threshold=6, med_nmne_threshold=6,
high_nmne_threshold=9, high_nmne_threshold=9,
include_nmne=True,
) )
assert nic_obs.high_nmne_threshold == 9 assert nic_obs.high_nmne_threshold == 9
@@ -85,6 +90,7 @@ def test_nic_categories(simulation):
low_nmne_threshold=9, low_nmne_threshold=9,
med_nmne_threshold=6, med_nmne_threshold=6,
high_nmne_threshold=9, high_nmne_threshold=9,
include_nmne=True,
) )
with pytest.raises(Exception): with pytest.raises(Exception):
@@ -94,4 +100,5 @@ def test_nic_categories(simulation):
low_nmne_threshold=3, low_nmne_threshold=3,
med_nmne_threshold=9, med_nmne_threshold=9,
high_nmne_threshold=9, high_nmne_threshold=9,
include_nmne=True,
) )

27
tests/integration_tests/game_layer/observations/test_node_observations.py
View File

@@ -19,15 +19,28 @@ def simulation(example_network) -> Simulation:
return sim return sim
def test_node_observation(simulation): def test_host_observation(simulation):
"""Test a Node observation.""" """Test a Host observation."""
pc: Computer = simulation.network.get_node_by_hostname("client_1") pc: Computer = simulation.network.get_node_by_hostname("client_1")
node_obs = HostObservation(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("operating_status") == 1 # computer is on
assert observation_state.get("SERVICES") is not None assert observation_state.get("SERVICES") is not None
@@ -36,11 +49,11 @@ def test_node_observation(simulation):
# turn off computer # turn off computer
pc.power_off() 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 assert observation_state.get("operating_status") == 4 # shutting down
for i in range(pc.shut_down_duration + 1): for i in range(pc.shut_down_duration + 1):
pc.apply_timestep(i) 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 assert observation_state.get("operating_status") == 2

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

@@ -14,7 +14,8 @@ def test_file_observation():
state = sim.describe_state() state = sim.describe_state()
dog_file_obs = FileObservation( 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.observe(state) == {"health_status": 1}
assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)}) assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)})

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

@@ -168,8 +168,8 @@ def test_capture_nmne_observations(uc2_network):
set_nmne_config(nmne_config) set_nmne_config(nmne_config)
# Define observations for the NICs of the database and web servers # Define observations for the NICs of the database and web servers
db_server_nic_obs = NICObservation(where=["network", "nodes", "database_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]) 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 # Iterate through a set of test cases to simulate multiple DELETE queries
for i in range(0, 20): for i in range(0, 20):