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

#902 - replaced 'final_node_<placeholder>' with 'reference_node_<placeholder>' in methods for scoring of os_state, file_system_state, service state and operating state. This fixed the reward function so it is checked at each step for node operating system state, operating state, file system state and service state.

Browse Source 
- Added unit tests.
This commit is contained in:
Chris McCarthy
2023-05-25 14:05:53 +01:00
parent 4cabc8a87a
commit ddb6adae2b
9 changed files with 526 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -37,6 +37,8 @@ from primaite.pol.ier import IER
from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
from primaite.transactions.transaction import Transaction
_LOGGER = logging.getLogger(__name__)
class Primaite(Env):
"""PRIMmary AI Training Evironment (Primaite) class."""
@@ -145,14 +147,12 @@ class Primaite(Env):
# Open the config file and build the environment laydown
try:
self.config_file = open(
"config/" + self.config_values.config_filename_use_case, "r"
)
self.config_file = open(self.config_values.config_filename_use_case, "r")
self.config_data = yaml.safe_load(self.config_file)
self.load_config()
except Exception:
logging.error("Could not load the environment configuration")
logging.error("Exception occured", exc_info=True)
_LOGGER.error("Could not load the environment configuration")
_LOGGER.error("Exception occured", exc_info=True)
# Store the node objects as node attributes
# (This is so we can access them as objects)
@@ -180,8 +180,8 @@ class Primaite(Env):
plt.savefig(filename, format="PNG")
plt.clf()
except Exception:
logging.error("Could not save network diagram")
logging.error("Exception occured", exc_info=True)
_LOGGER.error("Could not save network diagram")
_LOGGER.error("Exception occured", exc_info=True)
print("Could not save network diagram")
# Define Observation Space
@@ -223,7 +223,7 @@ class Primaite(Env):
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ACTION_TYPE.NODE:
logging.info("Action space type NODE selected")
_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, o/s state, service state, file system state)
@@ -238,7 +238,7 @@ class Primaite(Env):
]
)
else:
logging.info("Action space type ACL selected")
_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)
@@ -273,10 +273,10 @@ class Primaite(Env):
self.csv_writer = csv.writer(self.csv_file)
self.csv_writer.writerow(header)
except Exception:
logging.error(
_LOGGER.error(
"Could not create csv file to hold average reward per episode"
)
logging.error("Exception occured", exc_info=True)
_LOGGER.error("Exception occured", exc_info=True)
def reset(self):
"""
@@ -322,7 +322,7 @@ class Primaite(Env):
step_info: Additional information relating to this step
"""
if self.step_count == 0:
print("Episode: " + str(self.episode_count) + " running")
print(f"Episode: {str(self.episode_count)}")
# TEMP
done = False
@@ -402,7 +402,7 @@ class Primaite(Env):
self.step_count,
self.config_values,
)
# print("Step reward: " + str(reward))
print(f" Step {self.step_count} Reward: {str(reward)}")
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
@@ -410,7 +410,7 @@ class Primaite(Env):
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
print("Average reward: " + str(self.average_reward))
print(f" Average Reward: {str(self.average_reward)}")
# Load the reward into the transaction
transaction.set_reward(reward)
@@ -757,7 +757,7 @@ class Primaite(Env):
# Do nothing (bad formatting)
pass
logging.info("Environment configuration loaded")
_LOGGER.info("Environment configuration loaded")
print("Environment configuration loaded")
def create_node(self, item):
@@ -1090,7 +1090,7 @@ class Primaite(Env):
item: A config data item representing steps info
"""
self.episode_steps = int(steps_info["steps"])
logging.info("Training episodes have " + str(self.episode_steps) + " steps")
_LOGGER.info("Training episodes have " + str(self.episode_steps) + " steps")
def reset_environment(self):
"""

116
src/primaite/environment/reward.py
View File

@@ -98,27 +98,27 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_operating_state == HARDWARE_STATE.ON:
if final_node_operating_state == HARDWARE_STATE.OFF:
if reference_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_on
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_on
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.OFF:
if final_node_operating_state == HARDWARE_STATE.ON:
if reference_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_off
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_off
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
if final_node_operating_state == HARDWARE_STATE.ON:
if reference_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_resetting
elif final_node_operating_state == HARDWARE_STATE.OFF:
elif reference_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_resetting
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting
else:
pass
@@ -148,29 +148,29 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_os_state == SOFTWARE_STATE.GOOD:
if final_node_os_state == SOFTWARE_STATE.PATCHING:
if reference_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING:
if final_node_os_state == SOFTWARE_STATE.GOOD:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
elif final_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED:
if final_node_os_state == SOFTWARE_STATE.GOOD:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
else:
pass
@@ -204,46 +204,46 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
# Need to compare initial and reference state of node (i.e. at every step)
if initial_service.get_state() == SOFTWARE_STATE.GOOD:
if final_service.get_state() == SOFTWARE_STATE.PATCHING:
if reference_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_good
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed
else:
pass
@@ -276,67 +276,67 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
# Need to compare initial and reference state of node (i.e. at every step)
if initial_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing_should_be_good
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_good
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_good
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_good
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed
else:
pass

4
src/primaite/main.py
View File

@@ -165,7 +165,9 @@ def load_config_values():
config_values.agent_identifier = config_data["agentIdentifier"]
config_values.num_episodes = int(config_data["numEpisodes"])
config_values.time_delay = int(config_data["timeDelay"])
config_values.config_filename_use_case = config_data["configFilename"]
config_values.config_filename_use_case = (
"config/" + config_data["configFilename"]
)
config_values.session_type = config_data["sessionType"]
config_values.load_agent = bool(config_data["loadAgent"])
config_values.agent_load_file = config_data["agentLoadFile"]