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

901 - added changes back to ticket

Browse Source
This commit is contained in:
SunilSamra
2023-06-13 09:45:45 +01:00
parent cdd7183d85
commit cf64990cff
5 changed files with 50 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/acl/access_control_list.py
View File

@@ -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:

24
src/primaite/config/training_config.py
View File

@@ -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:

48
src/primaite/environment/primaite_env.py
View File

@@ -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

4
tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml
View File

@@ -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)

5
tests/config/obs_tests/main_config_without_obs.yaml
View File

@@ -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