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

#917 - Integrated both SB3 and RLlib agents into PrimaiteSession

Browse Source
This commit is contained in:
Chris McCarthy
2023-06-19 20:27:08 +01:00
parent c09874edbe
commit 3670f16766
13 changed files with 726 additions and 688 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/primaite/config/_package_data/training/training_config_main.yaml
View File

@@ -54,6 +54,13 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip]
# The high value for the observation space
observation_space_high_value: 1000000000
# The Agent output verbosity level:
# Options are:
# "NONE" (No Output)
# "INFO" (Info Messages)
# "ALL" (All Messages)
output_verbose_level: INFO
# Reward values
# Generic
all_ok: 0

56
src/primaite/config/lay_down_config.py
View File

@@ -1,21 +1,63 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from pathlib import Path
from typing import Final
from typing import Final, Union, Dict, Any
import networkx
import yaml
from primaite import USERS_CONFIG_DIR, getLogger
_LOGGER = getLogger(__name__)
_EXAMPLE_LAY_DOWN: Final[Path] = USERS_CONFIG_DIR / "example_config" / "lay_down"
_EXAMPLE_LAY_DOWN: Final[
Path] = USERS_CONFIG_DIR / "example_config" / "lay_down"
# class LayDownConfig:
# network: networkx.Graph
# POL
# EIR
# ACL
def convert_legacy_lay_down_config_dict(
legacy_config_dict: Dict[str, Any]
) -> Dict[str, Any]:
"""
Convert a legacy lay down config dict to the new format.
:param legacy_config_dict: A legacy lay down config dict.
"""
_LOGGER.warning("Legacy lay down config conversion not yet implemented")
return legacy_config_dict
def load(
file_path: Union[str, Path],
legacy_file: bool = False
) -> Dict:
"""
Read in a lay down config yaml file.
:param file_path: The config file path.
:param legacy_file: True if the config file is legacy format, otherwise
False.
:return: The lay down config as a dict.
:raises ValueError: If the file_path does not exist.
"""
if not isinstance(file_path, Path):
file_path = Path(file_path)
if file_path.exists():
with open(file_path, "r") as file:
config = yaml.safe_load(file)
_LOGGER.debug(f"Loading lay down config file: {file_path}")
if legacy_file:
try:
config = convert_legacy_lay_down_config_dict(config)
except KeyError:
msg = (
f"Failed to convert lay down config file {file_path} "
f"from legacy format. Attempting to use file as is."
)
_LOGGER.error(msg)
return config
msg = f"Cannot load the lay down config as it does not exist: {file_path}"
_LOGGER.error(msg)
raise ValueError(msg)
def ddos_basic_one_config_path() -> Path:
"""

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

@@ -10,11 +10,27 @@ import yaml
from primaite import USERS_CONFIG_DIR, getLogger
from primaite.common.enums import DeepLearningFramework
from primaite.common.enums import ActionType, RedAgentIdentifier, \
AgentFramework, SessionType
AgentFramework, SessionType, OutputVerboseLevel
_LOGGER = getLogger(__name__)
_EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training"
_EXAMPLE_TRAINING: Final[
Path] = USERS_CONFIG_DIR / "example_config" / "training"
def main_training_config_path() -> Path:
"""
The path to the example training_config_main.yaml file.
:return: The file path.
"""
path = _EXAMPLE_TRAINING / "training_config_main.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path
@dataclass()
@@ -24,44 +40,47 @@ class TrainingConfig:
"The AgentFramework"
deep_learning_framework: DeepLearningFramework = DeepLearningFramework.TF
"The DeepLearningFramework."
"The DeepLearningFramework"
red_agent_identifier: RedAgentIdentifier = RedAgentIdentifier.PPO
"The RedAgentIdentifier.."
"The RedAgentIdentifier"
action_type: ActionType = ActionType.ANY
"The ActionType to use."
"The ActionType to use"
num_episodes: int = 10
"The number of episodes to train over."
"The number of episodes to train over"
num_steps: int = 256
"The number of steps in an episode."
"The number of steps in an episode"
checkpoint_every_n_episodes: int = 5
"The agent will save a checkpoint every n episodes."
"The agent will save a checkpoint every n episodes"
observation_space: dict = field(
default_factory=lambda: {"components": [{"name": "NODE_LINK_TABLE"}]}
)
"The observation space config dict."
"The observation space config dict"
time_delay: int = 10
"The delay between steps (ms). Applies to generic agents only."
"The delay between steps (ms). Applies to generic agents only"
# file
session_type: SessionType = SessionType.TRAINING
"The type of PrimAITE session to run."
"The type of PrimAITE session to run"
load_agent: str = False
"Determine whether to load an agent from file."
"Determine whether to load an agent from file"
agent_load_file: Optional[str] = None
"File path and file name of agent if you're loading one in."
"File path and file name of agent if you're loading one in"
# Environment
observation_space_high_value: int = 1000000000
"The high value for the observation space."
"The high value for the observation space"
output_verbose_level: OutputVerboseLevel = OutputVerboseLevel.INFO
"The Agent output verbosity level"
# Reward values
# Generic
@@ -126,28 +145,28 @@ class TrainingConfig:
# Patching / Reset durations
os_patching_duration: int = 5
"The time taken to patch the OS."
"The time taken to patch the OS"
node_reset_duration: int = 5
"The time taken to reset a node (hardware)."
"The time taken to reset a node (hardware)"
node_booting_duration: int = 3
"The Time taken to turn on the node."
"The Time taken to turn on the node"
node_shutdown_duration: int = 2
"The time taken to turn off the node."
"The time taken to turn off the node"
service_patching_duration: int = 5
"The time taken to patch a service."
"The time taken to patch a service"
file_system_repairing_limit: int = 5
"The time take to repair the file system."
"The time take to repair the file system"
file_system_restoring_limit: int = 5
"The time take to restore the file system."
"The time take to restore the file system"
file_system_scanning_limit: int = 5
"The time taken to scan the file system."
"The time taken to scan the file system"
@classmethod
def from_dict(
@@ -157,9 +176,10 @@ class TrainingConfig:
field_enum_map = {
"agent_framework": AgentFramework,
"deep_learning_framework": DeepLearningFramework,
"red_agent_identifier": RedAgentIdentifier,
"action_type": ActionType,
"session_type": SessionType
"red_agent_identifier": RedAgentIdentifier,
"action_type": ActionType,
"session_type": SessionType,
"output_verbose_level": OutputVerboseLevel
}
for field, enum_class in field_enum_map.items():
@@ -178,28 +198,19 @@ class TrainingConfig:
"""
data = self.__dict__
if json_serializable:
data["agent_framework"] = self.agent_framework.value
data["deep_learning_framework"] = self.deep_learning_framework.value
data["red_agent_identifier"] = self.red_agent_identifier.value
data["action_type"] = self.action_type.value
data["output_verbose_level"] = self.output_verbose_level.value
return data
def main_training_config_path() -> Path:
"""
The path to the example training_config_main.yaml file.
:return: The file path.
"""
path = _EXAMPLE_TRAINING / "training_config_main.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
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.
@@ -246,7 +257,8 @@ def convert_legacy_training_config_dict(
agent_framework: AgentFramework = AgentFramework.SB3,
red_agent_identifier: RedAgentIdentifier = RedAgentIdentifier.PPO,
action_type: ActionType = ActionType.ANY,
num_steps: int = 256
num_steps: int = 256,
output_verbose_level: OutputVerboseLevel = OutputVerboseLevel.INFO
) -> Dict[str, Any]:
"""
Convert a legacy training config dict to the new format.
@@ -260,13 +272,16 @@ def convert_legacy_training_config_dict(
don't have action_type values.
:param num_steps: The number of steps to set as legacy training configs
don't have num_steps values.
:param output_verbose_level: The agent output verbose level to use as
legacy training configs don't have output_verbose_level values.
:return: The converted training config dict.
"""
config_dict = {
"agent_framework": agent_framework.name,
"red_agent_identifier": red_agent_identifier.name,
"action_type": action_type.name,
"num_steps": num_steps
"num_steps": num_steps,
"output_verbose_level": output_verbose_level
}
for legacy_key, value in legacy_config_dict.items():
new_key = _get_new_key_from_legacy(legacy_key)