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

Align observations to Common approach

Browse Source
This commit is contained in:
Marek Wolan
2023-10-24 11:07:25 +01:00
parent d4eee36b7b
commit 6b7c483a67
2 changed files with 59 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
example_config.yaml
View File

@@ -2,10 +2,10 @@ training_config:
rl_framework: SB3 rl_framework: SB3
rl_algorithm: PPO rl_algorithm: PPO
seed: 333 seed: 333
n_learn_episodes: 4 n_learn_episodes: 1
n_learn_steps: 128 n_learn_steps: 8
n_eval_episodes: 1 n_eval_episodes: 0
n_eval_steps: 128 n_eval_steps: 8
game_config: game_config:
@@ -39,10 +39,10 @@ game_config:
options: options:
nodes: nodes:
- node_ref: client_2 - node_ref: client_2
max_folders_per_node: 2 max_folders_per_node: 1
max_files_per_folder: 2 max_files_per_folder: 1
max_services_per_node: 2 max_services_per_node: 1
max_nics_per_node: 8 max_nics_per_node: 2
max_acl_rules: 10 max_acl_rules: 10
reward_function: reward_function:
@@ -93,9 +93,9 @@ game_config:
options: options:
nodes: nodes:
- node_ref: client_1 - node_ref: client_1
max_folders_per_node: 2 max_folders_per_node: 1
max_files_per_folder: 2 max_files_per_folder: 1
max_services_per_node: 2 max_services_per_node: 1
reward_function: reward_function:
reward_components: reward_components:
@@ -113,9 +113,10 @@ game_config:
observation_space: observation_space:
type: UC2BlueObservation type: UC2BlueObservation
options: options:
num_services_per_node: 2 num_services_per_node: 1
num_folders_per_node: 2 num_folders_per_node: 1
num_files_per_folder: 2 num_files_per_folder: 1
num_nics_per_node: 2
nodes: nodes:
- node_ref: domain_controller - node_ref: domain_controller
services: services:
@@ -148,6 +149,8 @@ game_config:
- link_ref: switch_2___client_2 - link_ref: switch_2___client_2
- link_ref: switch_2___security_suite - link_ref: switch_2___security_suite
acl: acl:
options:
max_acl_rules: 10
router_node_ref: router_1 router_node_ref: router_1
ip_address_order: ip_address_order:
- node_ref: domain_controller - node_ref: domain_controller

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

@@ -167,7 +167,7 @@ class ServiceObservation(AbstractObservation):
class LinkObservation(AbstractObservation): class LinkObservation(AbstractObservation):
"""Observation of a link in the network.""" """Observation of a link in the network."""
default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} default_observation: spaces.Space = {"PROTOCOLS": {"ALL": 0}}
"Default observation is what should be returned when the link doesn't exist." "Default observation is what should be returned when the link doesn't exist."
def __init__(self, where: Optional[Tuple[str]] = None) -> None: def __init__(self, where: Optional[Tuple[str]] = None) -> None:
@@ -206,7 +206,7 @@ class LinkObservation(AbstractObservation):
utilisation_category = int(utilisation_fraction * 10) + 1 utilisation_category = int(utilisation_fraction * 10) + 1
# TODO: once the links support separte load per protocol, this needs amendment to reflect that. # TODO: once the links support separte load per protocol, this needs amendment to reflect that.
return {"protocols": {"all": {"load": utilisation_category}}} return {"PROTOCOLS": {"ALL": utilisation_category}}
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
@@ -215,7 +215,7 @@ class LinkObservation(AbstractObservation):
:return: Gymnasium space :return: Gymnasium space
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})}) return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
@classmethod @classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "LinkObservation": def from_config(cls, config: Dict, session: "PrimaiteSession") -> "LinkObservation":
@@ -264,7 +264,6 @@ class FolderObservation(AbstractObservation):
truncated_file = self.files.pop() truncated_file = self.files.pop()
msg = f"Too many files in folde observation. Truncating file {truncated_file}" msg = f"Too many files in folde observation. Truncating file {truncated_file}"
_LOGGER.warn(msg) _LOGGER.warn(msg)
raise UserWarning(msg)
self.default_observation = { self.default_observation = {
"health_status": 0, "health_status": 0,
@@ -407,6 +406,7 @@ class NodeObservation(AbstractObservation):
num_services_per_node: int = 2, num_services_per_node: int = 2,
num_folders_per_node: int = 2, num_folders_per_node: int = 2,
num_files_per_folder: int = 2, num_files_per_folder: int = 2,
num_nics_per_node: int = 2,
) -> None: ) -> None:
""" """
Configurable observation for a node in the simulation. Configurable observation for a node in the simulation.
@@ -440,18 +440,25 @@ class NodeObservation(AbstractObservation):
truncated_service = self.services.pop() truncated_service = self.services.pop()
msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}" msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
_LOGGER.warn(msg) _LOGGER.warn(msg)
raise UserWarning(msg)
# truncate service list # truncate service list
self.folders: List[FolderObservation] = folders self.folders: List[FolderObservation] = folders
# add empty folder observation without `where` parameter that will always return default (blank) observations # add empty folder observation without `where` parameter that will always return default (blank) observations
while len(self.folders) < num_folders_per_node: while len(self.folders) < num_folders_per_node:
self.folders.append(FolderObservation()) self.folders.append(FolderObservation(num_files_per_folder=num_files_per_folder))
while len(self.folders) > num_folders_per_node: while len(self.folders) > num_folders_per_node:
truncated_folder = self.folders.pop() truncated_folder = self.folders.pop()
msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}" msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}"
_LOGGER.warn(msg)
self.nics: List[NicObservation] = nics self.nics: List[NicObservation] = nics
while len(self.nics) < num_nics_per_node:
self.nics.append(NicObservation())
while len(self.nics) > num_nics_per_node:
truncated_nic = self.nics.pop()
msg = f"Too many NICs in Node observation for node. Truncating service {truncated_nic.where[-1]}"
_LOGGER.warn(msg)
self.logon_status: bool = logon_status self.logon_status: bool = logon_status
self.default_observation: Dict = { self.default_observation: Dict = {
@@ -512,6 +519,7 @@ class NodeObservation(AbstractObservation):
num_services_per_node: int = 2, num_services_per_node: int = 2,
num_folders_per_node: int = 2, num_folders_per_node: int = 2,
num_files_per_folder: int = 2, num_files_per_folder: int = 2,
num_nics_per_node: int = 2,
) -> "NodeObservation": ) -> "NodeObservation":
"""Create node observation from a config. Also creates child service, folder and NIC observations. """Create node observation from a config. Also creates child service, folder and NIC observations.
@@ -562,6 +570,7 @@ class NodeObservation(AbstractObservation):
num_services_per_node=num_services_per_node, num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node, num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder, num_files_per_folder=num_files_per_folder,
num_nics_per_node=num_nics_per_node,
) )
@@ -605,7 +614,6 @@ class AclObservation(AbstractObservation):
self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)} self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)}
"List of protocols which are part of the game, defines ordering when converting to an ID" "List of protocols which are part of the game, defines ordering when converting to an ID"
self.default_observation: Dict = { self.default_observation: Dict = {
"RULES": {
i i
+ 1: { + 1: {
"position": i, "position": i,
@@ -618,7 +626,6 @@ class AclObservation(AbstractObservation):
} }
for i in range(self.num_rules) for i in range(self.num_rules)
} }
}
def observe(self, state: Dict) -> Dict: def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation. """Generate observation based on the current state of the simulation.
@@ -636,10 +643,9 @@ class AclObservation(AbstractObservation):
# TODO: what if the ACL has more rules than num of max rules for obs space # TODO: what if the ACL has more rules than num of max rules for obs space
obs = {} obs = {}
obs["RULES"] = {}
for i, rule_state in acl_state.items(): for i, rule_state in acl_state.items():
if rule_state is None: if rule_state is None:
obs["RULES"][i + 1] = { obs[i + 1] = {
"position": i, "position": i,
"permission": 0, "permission": 0,
"source_node_id": 0, "source_node_id": 0,
@@ -649,7 +655,7 @@ class AclObservation(AbstractObservation):
"protocol": 0, "protocol": 0,
} }
else: else:
obs["RULES"][i + 1] = { obs[i + 1] = {
"position": i, "position": i,
"permission": rule_state["action"], "permission": rule_state["action"],
"source_node_id": self.node_to_id[rule_state["src_ip_address"]], "source_node_id": self.node_to_id[rule_state["src_ip_address"]],
@@ -668,8 +674,6 @@ class AclObservation(AbstractObservation):
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict( return spaces.Dict(
{
"RULES": spaces.Dict(
{ {
i i
+ 1: spaces.Dict( + 1: spaces.Dict(
@@ -687,8 +691,6 @@ class AclObservation(AbstractObservation):
for i in range(self.num_rules) for i in range(self.num_rules)
} }
) )
}
)
@classmethod @classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation": def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation":
@@ -701,6 +703,7 @@ class AclObservation(AbstractObservation):
:return: Observation object :return: Observation object
:rtype: AclObservation :rtype: AclObservation
""" """
max_acl_rules = config["options"]["max_acl_rules"]
node_ip_to_idx = {} node_ip_to_idx = {}
for ip_idx, ip_map_config in enumerate(config["ip_address_order"]): for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
node_ref = ip_map_config["node_ref"] node_ref = ip_map_config["node_ref"]
@@ -715,6 +718,7 @@ class AclObservation(AbstractObservation):
ports=session.options.ports, ports=session.options.ports,
protocols=session.options.protocols, protocols=session.options.protocols,
where=["network", "nodes", router_uuid, "acl", "acl"], where=["network", "nodes", router_uuid, "acl", "acl"],
num_rules=max_acl_rules,
) )
@@ -846,6 +850,7 @@ class UC2BlueObservation(AbstractObservation):
num_services_per_node = config["num_services_per_node"] num_services_per_node = config["num_services_per_node"]
num_folders_per_node = config["num_folders_per_node"] num_folders_per_node = config["num_folders_per_node"]
num_files_per_folder = config["num_files_per_folder"] num_files_per_folder = config["num_files_per_folder"]
num_nics_per_node = config["num_nics_per_node"]
nodes = [ nodes = [
NodeObservation.from_config( NodeObservation.from_config(
config=n, config=n,
@@ -853,6 +858,7 @@ class UC2BlueObservation(AbstractObservation):
num_services_per_node=num_services_per_node, num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node, num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder, num_files_per_folder=num_files_per_folder,
num_nics_per_node=num_nics_per_node,
) )
for n in node_configs for n in node_configs
] ]