From cf64990cffd5bc1a46a82c71f5a638401ecf5a55 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 13 Jun 2023 09:45:45 +0100 Subject: [PATCH] 901 - added changes back to ticket --- src/primaite/acl/access_control_list.py | 2 +- src/primaite/config/training_config.py | 24 ++++++---- src/primaite/environment/primaite_env.py | 48 +++++++++++-------- .../main_config_ACCESS_CONTROL_LIST.yaml | 4 ++ .../obs_tests/main_config_without_obs.yaml | 5 +- 5 files changed, 50 insertions(+), 33 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 44a96743..d75b9756 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -97,7 +97,7 @@ class AccessControlList: _port: the port _position: position to insert ACL rule into ACL list (starting from index 1 and NOT 0) """ - position_index = int(_position) - 1 + position_index = int(_position) new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) if len(self.acl) < self.max_acl_rules: if len(self.acl) > position_index > -1: diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 4af36abe..67403c52 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -1,7 +1,7 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Final, Union, Optional +from typing import Any, Dict, Final, Optional, Union import yaml @@ -51,6 +51,16 @@ class TrainingConfig: observation_space_high_value: int = 1000000000 "The high value for the observation space." + # Access Control List/Rules + apply_implicit_rule: str = True + "User choice to have Implicit ALLOW or DENY." + + implicit_acl_rule: str = "DENY" + "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." + + max_number_acl_rule: int = 0 + "Sets a limit for number of acl rules allowed in the list and environment." + # Reward values # Generic all_ok: int = 0 @@ -167,8 +177,7 @@ def main_training_config_path() -> Path: return path -def load(file_path: Union[str, Path], - legacy_file: bool = False) -> TrainingConfig: +def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConfig: """ Read in a training config yaml file. @@ -213,9 +222,7 @@ def load(file_path: Union[str, Path], def convert_legacy_training_config_dict( - legacy_config_dict: Dict[str, Any], - num_steps: int = 256, - action_type: str = "ANY" + legacy_config_dict: Dict[str, Any], num_steps: int = 256, action_type: str = "ANY" ) -> Dict[str, Any]: """ Convert a legacy training config dict to the new format. @@ -227,10 +234,7 @@ def convert_legacy_training_config_dict( don't have action_type values. :return: The converted training config dict. """ - config_dict = { - "num_steps": num_steps, - "action_type": action_type - } + config_dict = {"num_steps": num_steps, "action_type": action_type} for legacy_key, value in legacy_config_dict.items(): new_key = _get_new_key_from_legacy(legacy_key) if new_key: diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index cd0c660e..0a351b08 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -118,7 +118,10 @@ class Primaite(Env): self.red_node_pol = {} # Create the Access Control List - self.acl = AccessControlList() + self.acl = AccessControlList( + self.training_config.implicit_acl_rule, + self.training_config.max_number_acl_rule, + ) # Create a list of services (enums) self.services_list = [] @@ -212,22 +215,10 @@ class Primaite(Env): # Define Action Space - depends on action space type (Node or ACL) if self.training_config.action_type == ActionType.NODE: _LOGGER.info("Action space type NODE selected") - # Terms (for node action space): - # [0, num nodes] - node ID (0 = nothing, node ID) - # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa - # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa - # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa self.action_dict = self.create_node_action_dict() self.action_space = spaces.Discrete(len(self.action_dict)) elif self.training_config.action_type == ActionType.ACL: _LOGGER.info("Action space type ACL selected") - # Terms (for ACL action space): - # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - # [0, 1] - Permission (0 = DENY, 1 = ALLOW) - # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) self.action_dict = self.create_acl_action_dict() self.action_space = spaces.Discrete(len(self.action_dict)) elif self.training_config.action_type == ActionType.ANY: @@ -1144,6 +1135,11 @@ class Primaite(Env): } """ + # Terms (for node action space): + # [0, num nodes] - node ID (0 = nothing, node ID) + # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa + # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa + # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa # reserve 0 action to be a nothing action actions = {0: [1, 0, 0, 0]} action_key = 1 @@ -1165,6 +1161,14 @@ class Primaite(Env): def create_acl_action_dict(self): """Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.""" + # Terms (for ACL action space): + # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) + # [0, 1] - Permission (0 = DENY, 1 = ALLOW) + # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) + # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) + # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) # reserve 0 action to be a nothing action actions = {0: [0, 0, 0, 0, 0, 0]} @@ -1178,14 +1182,16 @@ class Primaite(Env): for dest_ip in range(self.num_nodes + 1): for protocol in range(self.num_services + 1): for port in range(self.num_ports + 1): - action = [ - action_decision, - action_permission, - source_ip, - dest_ip, - protocol, - port, - ] + for position in range(self.max_acl_rules - 1): + action = [ + action_decision, + action_permission, + source_ip, + dest_ip, + protocol, + port, + position, + ] # Check to see if its an action we want to include as possible i.e. not a nothing action if is_valid_acl_action_extra(action): actions[action_key] = action diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index b36cd6ce..856e963d 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -14,7 +14,11 @@ observationSpace: implicit_acl_rule: DENY max_number_of_acl_rules: 10 +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment max_number_acl_rules: 10 numEpisodes: 1 # Time delay between steps (for generic agents) diff --git a/tests/config/obs_tests/main_config_without_obs.yaml b/tests/config/obs_tests/main_config_without_obs.yaml index 99005678..c671b31f 100644 --- a/tests/config/obs_tests/main_config_without_obs.yaml +++ b/tests/config/obs_tests/main_config_without_obs.yaml @@ -27,8 +27,11 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] # Environment config values # The high value for the observation space observation_space_high_value: 1_000_000_000 - +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment max_number_acl_rules: 10 # Reward values # Generic