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

Ran pre-commit hook on all files and performed changes to fix flake8 failures

Browse Source
This commit is contained in:
Chris McCarthy
2023-05-25 11:42:19 +01:00
parent aa8284897a
commit 4eb1658966
42 changed files with 1371 additions and 1170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

341
src/primaite/main.py
View File

@@ -1,37 +1,31 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Primaite - main (harness) module
Primaite - main (harness) module.
Coding Standards: PEP 8
"""
from sys import exc_info
import time
import yaml
import os.path
import logging
import os.path
import time
from datetime import datetime
import yaml
from stable_baselines3 import A2C, PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.ppo import MlpPolicy as PPOMlp
from primaite.common.config_values_main import config_values_main
from primaite.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import write_transaction_to_file
from primaite.common.config_values_main import config_values_main
from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy as PPOMlp
from stable_baselines3 import A2C
from stable_baselines3.common.env_checker import check_env
from stable_baselines3.common.evaluation import evaluate_policy
# FUNCTIONS #
################################# FUNCTIONS ######################################
def run_generic():
"""
Run against a generic agent
"""
"""Run against a generic agent."""
for episode in range(0, config_values.num_episodes):
for step in range(0, config_values.num_steps):
# Send the observation space to the agent to get an action
# TEMP - random action for now
# action = env.blue_agent_action(obs)
@@ -54,15 +48,20 @@ def run_generic():
def run_stable_baselines3_ppo():
"""
Run against a stable_baselines3 PPO agent
"""
"""Run against a stable_baselines3 PPO agent."""
if config_values.load_agent == True:
try:
agent = PPO.load(config_values.agent_load_file, env, verbose=0, n_steps=config_values.num_steps)
except:
print("ERROR: Could not load agent at location: " + config_values.agent_load_file)
agent = PPO.load(
config_values.agent_load_file,
env,
verbose=0,
n_steps=config_values.num_steps,
)
except Exception:
print(
"ERROR: Could not load agent at location: "
+ config_values.agent_load_file
)
logging.error("Could not load agent")
logging.error("Exception occured", exc_info=True)
else:
@@ -73,7 +72,7 @@ def run_stable_baselines3_ppo():
print("Starting training session...")
logging.info("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
agent.learn(total_timesteps=1)
save_agent(agent)
else:
# Default to being in an evaluation session
@@ -83,16 +82,22 @@ def run_stable_baselines3_ppo():
env.close()
def run_stable_baselines3_a2c():
"""
Run against a stable_baselines3 A2C agent
"""
def run_stable_baselines3_a2c():
"""Run against a stable_baselines3 A2C agent."""
if config_values.load_agent == True:
try:
agent = A2C.load(config_values.agent_load_file, env, verbose=0, n_steps=config_values.num_steps)
except:
print("ERROR: Could not load agent at location: " + config_values.agent_load_file)
agent = A2C.load(
config_values.agent_load_file,
env,
verbose=0,
n_steps=config_values.num_steps,
)
except Exception:
print(
"ERROR: Could not load agent at location: "
+ config_values.agent_load_file
)
logging.error("Could not load agent")
logging.error("Exception occured", exc_info=True)
else:
@@ -103,143 +108,213 @@ def run_stable_baselines3_a2c():
print("Starting training session...")
logging.info("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
agent.learn(total_timesteps=1)
save_agent(agent)
else:
# Default to being in an evaluation session
print("Starting evaluation session...")
logging.info("Starting evaluation session...")
evaluate_policy(agent, env, n_eval_episodes=config_values.num_episodes)
evaluate_policy(agent, env, n_eval_episodes=config_values.num_episodes)
env.close()
def save_agent(_agent):
"""
Persist an agent (only works for stable baselines3 agents at present)
"""
now = datetime.now() # current date and time
def save_agent(_agent):
"""Persist an agent (only works for stable baselines3 agents at present)."""
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
try:
path = 'outputs/agents/'
path = "outputs/agents/"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/agents/agent_saved_" + time
_agent.save(filename)
logging.info("Trained agent saved as " + filename)
except Exception as e:
except Exception:
logging.error("Could not save agent")
logging.error("Exception occured", exc_info=True)
def configure_logging():
"""
Configures logging
"""
def configure_logging():
"""Configures logging."""
try:
now = datetime.now() # current date and time
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
filename = "logs/app_" + time + ".log"
path = 'logs/'
path = "logs/"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
logging.basicConfig(filename=filename, filemode='w', format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)
except:
logging.basicConfig(
filename=filename,
filemode="w",
format="%(asctime)s - %(levelname)s - %(message)s",
datefmt="%d-%b-%y %H:%M:%S",
level=logging.INFO,
)
except Exception:
print("ERROR: Could not start logging")
def load_config_values():
"""
Loads the config values from the main config file into a config object
"""
def load_config_values():
"""Loads the config values from the main config file into a config object."""
try:
# Generic
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.session_type = config_data['sessionType']
config_values.load_agent = bool(config_data['loadAgent'])
config_values.agent_load_file = config_data['agentLoadFile']
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.session_type = config_data["sessionType"]
config_values.load_agent = bool(config_data["loadAgent"])
config_values.agent_load_file = config_data["agentLoadFile"]
# Environment
config_values.observation_space_high_value = int(config_data['observationSpaceHighValue'])
config_values.observation_space_high_value = int(
config_data["observationSpaceHighValue"]
)
# Reward values
# Generic
config_values.all_ok = int(config_data['allOk'])
config_values.all_ok = int(config_data["allOk"])
# Node Operating State
config_values.off_should_be_on = int(config_data['offShouldBeOn'])
config_values.off_should_be_resetting = int(config_data['offShouldBeResetting'])
config_values.on_should_be_off = int(config_data['onShouldBeOff'])
config_values.on_should_be_resetting = int(config_data['onShouldBeResetting'])
config_values.resetting_should_be_on = int(config_data['resettingShouldBeOn'])
config_values.resetting_should_be_off = int(config_data['resettingShouldBeOff'])
config_values.resetting = int(config_data['resetting'])
config_values.off_should_be_on = int(config_data["offShouldBeOn"])
config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
config_values.on_should_be_off = int(config_data["onShouldBeOff"])
config_values.on_should_be_resetting = int(config_data["onShouldBeResetting"])
config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"])
config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
config_values.resetting = int(config_data["resetting"])
# Node O/S or Service State
config_values.good_should_be_patching = int(config_data['goodShouldBePatching'])
config_values.good_should_be_compromised = int(config_data['goodShouldBeCompromised'])
config_values.good_should_be_overwhelmed = int(config_data['goodShouldBeOverwhelmed'])
config_values.patching_should_be_good = int(config_data['patchingShouldBeGood'])
config_values.patching_should_be_compromised = int(config_data['patchingShouldBeCompromised'])
config_values.patching_should_be_overwhelmed = int(config_data['patchingShouldBeOverwhelmed'])
config_values.patching = int(config_data['patching'])
config_values.compromised_should_be_good = int(config_data['compromisedShouldBeGood'])
config_values.compromised_should_be_patching = int(config_data['compromisedShouldBePatching'])
config_values.compromised_should_be_overwhelmed = int(config_data['compromisedShouldBeOverwhelmed'])
config_values.compromised = int(config_data['compromised'])
config_values.overwhelmed_should_be_good = int(config_data['overwhelmedShouldBeGood'])
config_values.overwhelmed_should_be_patching = int(config_data['overwhelmedShouldBePatching'])
config_values.overwhelmed_should_be_compromised = int(config_data['overwhelmedShouldBeCompromised'])
config_values.overwhelmed = int(config_data['overwhelmed'])
config_values.good_should_be_patching = int(config_data["goodShouldBePatching"])
config_values.good_should_be_compromised = int(
config_data["goodShouldBeCompromised"]
)
config_values.good_should_be_overwhelmed = int(
config_data["goodShouldBeOverwhelmed"]
)
config_values.patching_should_be_good = int(config_data["patchingShouldBeGood"])
config_values.patching_should_be_compromised = int(
config_data["patchingShouldBeCompromised"]
)
config_values.patching_should_be_overwhelmed = int(
config_data["patchingShouldBeOverwhelmed"]
)
config_values.patching = int(config_data["patching"])
config_values.compromised_should_be_good = int(
config_data["compromisedShouldBeGood"]
)
config_values.compromised_should_be_patching = int(
config_data["compromisedShouldBePatching"]
)
config_values.compromised_should_be_overwhelmed = int(
config_data["compromisedShouldBeOverwhelmed"]
)
config_values.compromised = int(config_data["compromised"])
config_values.overwhelmed_should_be_good = int(
config_data["overwhelmedShouldBeGood"]
)
config_values.overwhelmed_should_be_patching = int(
config_data["overwhelmedShouldBePatching"]
)
config_values.overwhelmed_should_be_compromised = int(
config_data["overwhelmedShouldBeCompromised"]
)
config_values.overwhelmed = int(config_data["overwhelmed"])
# Node File System State
config_values.good_should_be_repairing = int(config_data['goodShouldBeRepairing'])
config_values.good_should_be_restoring = int(config_data['goodShouldBeRestoring'])
config_values.good_should_be_corrupt = int(config_data['goodShouldBeCorrupt'])
config_values.good_should_be_destroyed = int(config_data['goodShouldBeDestroyed'])
config_values.repairing_should_be_good = int(config_data['repairingShouldBeGood'])
config_values.repairing_should_be_restoring = int(config_data['repairingShouldBeRestoring'])
config_values.repairing_should_be_corrupt = int(config_data['repairingShouldBeCorrupt'])
config_values.repairing_should_be_destroyed = int(config_data['repairingShouldBeDestroyed'])
config_values.repairing = int(config_data['repairing'])
config_values.restoring_should_be_good = int(config_data['restoringShouldBeGood'])
config_values.restoring_should_be_repairing = int(config_data['restoringShouldBeRepairing'])
config_values.restoring_should_be_corrupt = int(config_data['restoringShouldBeCorrupt'])
config_values.restoring_should_be_destroyed = int(config_data['restoringShouldBeDestroyed'])
config_values.restoring = int(config_data['restoring'])
config_values.corrupt_should_be_good = int(config_data['corruptShouldBeGood'])
config_values.corrupt_should_be_repairing = int(config_data['corruptShouldBeRepairing'])
config_values.corrupt_should_be_restoring = int(config_data['corruptShouldBeRestoring'])
config_values.corrupt_should_be_destroyed = int(config_data['corruptShouldBeDestroyed'])
config_values.corrupt = int(config_data['corrupt'])
config_values.destroyed_should_be_good = int(config_data['destroyedShouldBeGood'])
config_values.destroyed_should_be_repairing = int(config_data['destroyedShouldBeRepairing'])
config_values.destroyed_should_be_restoring = int(config_data['destroyedShouldBeRestoring'])
config_values.destroyed_should_be_corrupt = int(config_data['destroyedShouldBeCorrupt'])
config_values.destroyed = int(config_data['destroyed'])
config_values.scanning = int(config_data['scanning'])
config_values.good_should_be_repairing = int(
config_data["goodShouldBeRepairing"]
)
config_values.good_should_be_restoring = int(
config_data["goodShouldBeRestoring"]
)
config_values.good_should_be_corrupt = int(config_data["goodShouldBeCorrupt"])
config_values.good_should_be_destroyed = int(
config_data["goodShouldBeDestroyed"]
)
config_values.repairing_should_be_good = int(
config_data["repairingShouldBeGood"]
)
config_values.repairing_should_be_restoring = int(
config_data["repairingShouldBeRestoring"]
)
config_values.repairing_should_be_corrupt = int(
config_data["repairingShouldBeCorrupt"]
)
config_values.repairing_should_be_destroyed = int(
config_data["repairingShouldBeDestroyed"]
)
config_values.repairing = int(config_data["repairing"])
config_values.restoring_should_be_good = int(
config_data["restoringShouldBeGood"]
)
config_values.restoring_should_be_repairing = int(
config_data["restoringShouldBeRepairing"]
)
config_values.restoring_should_be_corrupt = int(
config_data["restoringShouldBeCorrupt"]
)
config_values.restoring_should_be_destroyed = int(
config_data["restoringShouldBeDestroyed"]
)
config_values.restoring = int(config_data["restoring"])
config_values.corrupt_should_be_good = int(config_data["corruptShouldBeGood"])
config_values.corrupt_should_be_repairing = int(
config_data["corruptShouldBeRepairing"]
)
config_values.corrupt_should_be_restoring = int(
config_data["corruptShouldBeRestoring"]
)
config_values.corrupt_should_be_destroyed = int(
config_data["corruptShouldBeDestroyed"]
)
config_values.corrupt = int(config_data["corrupt"])
config_values.destroyed_should_be_good = int(
config_data["destroyedShouldBeGood"]
)
config_values.destroyed_should_be_repairing = int(
config_data["destroyedShouldBeRepairing"]
)
config_values.destroyed_should_be_restoring = int(
config_data["destroyedShouldBeRestoring"]
)
config_values.destroyed_should_be_corrupt = int(
config_data["destroyedShouldBeCorrupt"]
)
config_values.destroyed = int(config_data["destroyed"])
config_values.scanning = int(config_data["scanning"])
# IER status
config_values.red_ier_running = int(config_data['redIerRunning'])
config_values.green_ier_blocked = int(config_data['greenIerBlocked'])
config_values.red_ier_running = int(config_data["redIerRunning"])
config_values.green_ier_blocked = int(config_data["greenIerBlocked"])
# Patching / Reset durations
config_values.os_patching_duration = int(config_data['osPatchingDuration'])
config_values.node_reset_duration = int(config_data['nodeResetDuration'])
config_values.service_patching_duration = int(config_data['servicePatchingDuration'])
config_values.file_system_repairing_limit = int(config_data['fileSystemRepairingLimit'])
config_values.file_system_restoring_limit = int(config_data['fileSystemRestoringLimit'])
config_values.file_system_scanning_limit = int(config_data['fileSystemScanningLimit'])
logging.info("Training agent: " + config_values.agent_identifier)
logging.info("Training environment config: " + config_values.config_filename_use_case)
logging.info("Training cycle has " + str(config_values.num_episodes) + " episodes")
config_values.os_patching_duration = int(config_data["osPatchingDuration"])
config_values.node_reset_duration = int(config_data["nodeResetDuration"])
config_values.service_patching_duration = int(
config_data["servicePatchingDuration"]
)
config_values.file_system_repairing_limit = int(
config_data["fileSystemRepairingLimit"]
)
config_values.file_system_restoring_limit = int(
config_data["fileSystemRestoringLimit"]
)
config_values.file_system_scanning_limit = int(
config_data["fileSystemScanningLimit"]
)
except Exception as e:
logging.info("Training agent: " + config_values.agent_identifier)
logging.info(
"Training environment config: " + config_values.config_filename_use_case
)
logging.info(
"Training cycle has " + str(config_values.num_episodes) + " episodes"
)
except Exception:
logging.error("Could not save load config data")
logging.error("Exception occured", exc_info=True)
################################# MAIN PROCESS ############################################
# MAIN PROCESS #
# Starting point
@@ -257,25 +332,25 @@ try:
config_values = config_values_main()
# Load in config data
load_config_values()
except Exception as e:
except Exception:
logging.error("Could not load main config")
logging.error("Exception occured", exc_info=True)
# Create a list of transactions
# A transaction is an object holding the:
# - episode #
# - step #
# A transaction is an object holding the:
# - episode #
# - step #
# - initial observation space
# - action
# - reward
# - new observation space
# - new observation space
transaction_list = []
# Create the Primaite environment
try:
env = Primaite(config_values, transaction_list)
logging.info("PrimAITE environment created")
except Exception as e:
except Exception:
logging.error("Could not create PrimAITE environment")
logging.error("Exception occured", exc_info=True)
@@ -302,11 +377,3 @@ config_file_main.close
print("Finished")
logging.info("Finished")