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

1429 - added code from ADSP branch to primaite_env.py and added NONE = 0 to NodePOLType in enums.py

Browse Source
This commit is contained in:
SunilSamra
2023-05-26 14:29:02 +01:00
parent 9d3d8d5945
commit dc011a489c
2 changed files with 27 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/primaite/common/enums.py
View File

@@ -49,6 +49,7 @@ class SoftwareState(Enum):
class NodePOLType(Enum): class NodePOLType(Enum):
"""Node Pattern of Life type enumeration.""" """Node Pattern of Life type enumeration."""
NONE = 0
OPERATING = 1 OPERATING = 1
OS = 2 OS = 2
SERVICE = 3 SERVICE = 3

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

@@ -247,7 +247,7 @@ class Primaite(Env):
self.action_dict = self.create_acl_action_dict() self.action_dict = self.create_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict)) self.action_space = spaces.Discrete(len(self.action_dict))
else: else:
_LOGGER.info("Action space type ANY selected") _LOGGER.info("Action space type ANY selected - Node + ACL")
self.action_dict = self.create_node_and_acl_action_dict() self.action_dict = self.create_node_and_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict)) self.action_space = spaces.Discrete(len(self.action_dict))
@@ -453,8 +453,18 @@ class Primaite(Env):
# At the moment, actions are only affecting nodes # At the moment, actions are only affecting nodes
if self.action_type == ActionType.NODE: if self.action_type == ActionType.NODE:
self.apply_actions_to_nodes(_action) self.apply_actions_to_nodes(_action)
else: elif self.action_type == ActionType.ACL:
self.apply_actions_to_acl(_action) self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 6
): # ACL actions in multidiscrete form have len 6
self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 4
): # Node actions in multdiscrete (array) from have len 4
self.apply_actions_to_nodes(_action)
else:
logging.error("Invalid action type found")
def apply_actions_to_nodes(self, _action): def apply_actions_to_nodes(self, _action):
""" """
@@ -463,10 +473,11 @@ class Primaite(Env):
Args: Args:
_action: The action space from the agent _action: The action space from the agent
""" """
node_id = _action[0] readable_action = self.action_dict[_action]
node_property = _action[1] node_id = readable_action[0]
property_action = _action[2] node_property = readable_action[1]
service_index = _action[3] property_action = readable_action[2]
service_index = readable_action[3]
# Check that the action is requesting a valid node # Check that the action is requesting a valid node
try: try:
@@ -552,12 +563,15 @@ class Primaite(Env):
Args: Args:
_action: The action space from the agent _action: The action space from the agent
""" """
action_decision = _action[0] # Convert discrete value back to multidiscrete
action_permission = _action[1] multidiscrete_action = self.action_dict[_action]
action_source_ip = _action[2]
action_destination_ip = _action[3] action_decision = multidiscrete_action[0]
action_protocol = _action[4] action_permission = multidiscrete_action[1]
action_port = _action[5] action_source_ip = multidiscrete_action[2]
action_destination_ip = multidiscrete_action[3]
action_protocol = multidiscrete_action[4]
action_port = multidiscrete_action[5]
if action_decision == 0: if action_decision == 0:
# It's decided to do nothing # It's decided to do nothing