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

Merge pull request #5 from Autonomous-Resilient-Cyber-Defence/dev

Browse Source 
Finalised the support for legacy config files
This commit is contained in:
Chris McCarthy
2023-07-28 14:19:15 +01:00
committed by GitHub
parent e647b35f6f 10ce2923c7
commit 5893ea8db4
16 changed files with 1454 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -1,5 +1,7 @@
# PrimAITE
![image](https://github.com/Autonomous-Resilient-Cyber-Defence/PrimAITE/assets/107395948/fdefa884-1105-44da-88fe-e3a1c98ee361)
The ARCD Primary-level AI Training Environment (**PrimAITE**) provides an effective simulation capability for the purposes of training and evaluating AI in a cyber-defensive role. It incorporates the functionality required of a primary-level ARCD environment, which includes:
- The ability to model a relevant platform / system context;

14
src/primaite/agents/agent_abc.py
View File

@@ -52,6 +52,8 @@ class AgentSessionABC(ABC):
training_config_path: Optional[Union[str, Path]] = None,
lay_down_config_path: Optional[Union[str, Path]] = None,
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Initialise an agent session from config files, or load a previous session.
@@ -64,6 +66,10 @@ class AgentSessionABC(ABC):
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param session_path: directory path of the session to load
"""
# initialise variables
@@ -72,6 +78,8 @@ class AgentSessionABC(ABC):
self._can_learn: bool = False
self._can_evaluate: bool = False
self.is_eval = False
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
self.session_timestamp: datetime = datetime.now()
@@ -91,12 +99,14 @@ class AgentSessionABC(ABC):
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Union[Path, str] = training_config_path
self._training_config: TrainingConfig = training_config.load(self._training_config_path)
self._training_config: TrainingConfig = training_config.load(
self._training_config_path, legacy_file=legacy_training_config
)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Union[Path, str] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config)
self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
# set random UUID for session

12
src/primaite/agents/sb3.py
View File

@@ -26,6 +26,8 @@ class SB3Agent(AgentSessionABC):
training_config_path: Optional[Union[str, Path]] = None,
lay_down_config_path: Optional[Union[str, Path]] = None,
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Initialise the SB3 Agent training session.
@@ -35,11 +37,17 @@ class SB3Agent(AgentSessionABC):
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "SB3")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
super().__init__(training_config_path, lay_down_config_path, session_path)
super().__init__(
training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config
)
if not self._training_config.agent_framework == AgentFramework.SB3:
msg = f"Expected SB3 agent_framework, " f"got {self._training_config.agent_framework}"
_LOGGER.error(msg)
@@ -75,6 +83,8 @@ class SB3Agent(AgentSessionABC):
lay_down_config_path=self._lay_down_config_path,
session_path=self.session_path,
timestamp_str=self.timestamp_str,
legacy_training_config=self.legacy_training_config,
legacy_lay_down_config=self.legacy_lay_down_config,
)
# check if there is a zip file that needs to be loaded

23
src/primaite/cli.py
View File

@@ -18,9 +18,9 @@ app = typer.Typer()
@app.command()
def build_dirs() -> None:
"""Build the PrimAITE app directories."""
from primaite.setup import setup_app_dirs
from primaite import PRIMAITE_PATHS
setup_app_dirs.run()
PRIMAITE_PATHS.mkdirs()
@app.command()
@@ -137,7 +137,13 @@ def setup(overwrite_existing: bool = True) -> None:
@app.command()
def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[str] = None) -> None:
def session(
tc: Optional[str] = None,
ldc: Optional[str] = None,
load: Optional[str] = None,
legacy_tc: bool = False,
legacy_ldc: bool = False,
) -> None:
"""
Run a PrimAITE session.
@@ -153,6 +159,10 @@ def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[
will use the default training config and laydown config. Inversely, if a training config and laydown config
is passed while a session directory is passed, PrimAITE will load the session and ignore the training config
and laydown config.
legacy_tc: If the training config file is a legacy file from PrimAITE < 2.0.
legacy_ldf: If the lay down config file is a legacy file from PrimAITE < 2.0.
"""
from primaite.config.lay_down_config import dos_very_basic_config_path
from primaite.config.training_config import main_training_config_path
@@ -170,7 +180,12 @@ def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[
if not ldc:
ldc = dos_very_basic_config_path()
run(training_config_path=tc, lay_down_config_path=ldc)
run(
training_config_path=tc,
lay_down_config_path=ldc,
legacy_training_config=legacy_tc,
legacy_lay_down_config=legacy_ldc,
)
@app.command()

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

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
from logging import Logger
from pathlib import Path
from typing import Any, Dict, Final, Union
from typing import Any, Dict, Final, List, Union
import yaml
@@ -12,14 +12,43 @@ _LOGGER: Logger = getLogger(__name__)
_EXAMPLE_LAY_DOWN: Final[Path] = PRIMAITE_PATHS.user_config_path / "example_config" / "lay_down"
def convert_legacy_lay_down_config_dict(legacy_config_dict: Dict[str, Any]) -> Dict[str, Any]:
def convert_legacy_lay_down_config(legacy_config: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Convert a legacy lay down config dict to the new format.
Convert a legacy lay down config to the new format.
:param legacy_config_dict: A legacy lay down config dict.
:param legacy_config: A legacy lay down config.
"""
_LOGGER.warning("Legacy lay down config conversion not yet implemented")
return legacy_config_dict
field_conversion_map = {
"itemType": "item_type",
"portsList": "ports_list",
"serviceList": "service_list",
"baseType": "node_class",
"nodeType": "node_type",
"hardwareState": "hardware_state",
"softwareState": "software_state",
"startStep": "start_step",
"endStep": "end_step",
"fileSystemState": "file_system_state",
"ipAddress": "ip_address",
"missionCriticality": "mission_criticality",
}
new_config = []
for item in legacy_config:
if "itemType" in item:
if item["itemType"] in ["ACTIONS", "STEPS"]:
continue
new_dict = {}
for key in item.keys():
conversion_key = field_conversion_map.get(key)
if key == "id" and "itemType" in item:
if item["itemType"] == "NODE":
conversion_key = "node_id"
if conversion_key:
new_dict[conversion_key] = item[key]
else:
new_dict[key] = item[key]
new_config.append(new_dict)
return new_config
def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
@@ -39,7 +68,7 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
_LOGGER.debug(f"Loading lay down config file: {file_path}")
if legacy_file:
try:
config = convert_legacy_lay_down_config_dict(config)
config = convert_legacy_lay_down_config(config)
except KeyError:
msg = (
f"Failed to convert lay down config file {file_path} "

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

@@ -291,12 +291,14 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConf
if legacy_file:
try:
config = convert_legacy_training_config_dict(config)
except KeyError:
except KeyError as e:
msg = (
f"Failed to convert training config file {file_path} "
f"from legacy format. Attempting to use file as is."
)
_LOGGER.error(msg)
raise e
try:
return TrainingConfig.from_dict(config)
except TypeError as e:
@@ -314,6 +316,9 @@ def convert_legacy_training_config_dict(
agent_identifier: AgentIdentifier = AgentIdentifier.PPO,
action_type: ActionType = ActionType.ANY,
num_train_steps: int = 256,
num_eval_steps: int = 256,
num_train_episodes: int = 10,
num_eval_episodes: int = 1,
) -> Dict[str, Any]:
"""
Convert a legacy training config dict to the new format.
@@ -325,8 +330,14 @@ def convert_legacy_training_config_dict(
training configs don't have agent_identifier values.
:param action_type: The action space type to set as legacy training configs
don't have action_type values.
:param num_train_steps: The number of steps to set as legacy training configs
:param num_train_steps: The number of train steps to set as legacy training configs
don't have num_train_steps values.
:param num_eval_steps: The number of eval steps to set as legacy training configs
don't have num_eval_steps values.
:param num_train_episodes: The number of train episodes to set as legacy training configs
don't have num_train_episodes values.
:param num_eval_episodes: The number of eval episodes to set as legacy training configs
don't have num_eval_episodes values.
:return: The converted training config dict.
"""
config_dict = {
@@ -334,6 +345,9 @@ def convert_legacy_training_config_dict(
"agent_identifier": agent_identifier.name,
"action_type": action_type.name,
"num_train_steps": num_train_steps,
"num_eval_steps": num_eval_steps,
"num_train_episodes": num_train_episodes,
"num_eval_episodes": num_eval_episodes,
"sb3_output_verbose_level": SB3OutputVerboseLevel.INFO.name,
}
session_type_map = {"TRAINING": "TRAIN", "EVALUATION": "EVAL"}

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

@@ -10,7 +10,6 @@ from typing import Any, Dict, Final, List, Tuple, Union
import networkx as nx
import numpy as np
import yaml
from gym import Env, spaces
from matplotlib import pyplot as plt
@@ -34,6 +33,7 @@ from primaite.common.enums import (
)
from primaite.common.service import Service
from primaite.config import training_config
from primaite.config.lay_down_config import load
from primaite.config.training_config import TrainingConfig
from primaite.environment.observations import ObservationsHandler
from primaite.environment.reward import calculate_reward_function
@@ -68,6 +68,8 @@ class Primaite(Env):
lay_down_config_path: Union[str, Path],
session_path: Path,
timestamp_str: str,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
The Primaite constructor.
@@ -76,13 +78,19 @@ class Primaite(Env):
:param lay_down_config_path: The lay down config filepath.
:param session_path: The directory path the session is writing to.
:param timestamp_str: The session timestamp in the format: <yyyy-mm-dd>_<hh-mm- ss>.
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
self.session_path: Final[Path] = session_path
self.timestamp_str: Final[str] = timestamp_str
self._training_config_path: Union[str, Path] = training_config_path
self._lay_down_config_path: Union[str, Path] = lay_down_config_path
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
self.training_config: TrainingConfig = training_config.load(training_config_path)
self.training_config: TrainingConfig = training_config.load(training_config_path, self.legacy_training_config)
_LOGGER.info(f"Using: {str(self.training_config)}")
# Number of steps in an episode
@@ -191,11 +199,8 @@ class Primaite(Env):
self._obs_space_description: List[str] = None
"The env observation space description for transactions writing"
# Open the config file and build the environment laydown
with open(self._lay_down_config_path, "r") as file:
# Open the config file and build the environment laydown
self.lay_down_config = yaml.safe_load(file)
self.load_lay_down_config()
self.lay_down_config = load(self._lay_down_config_path, self.legacy_lay_down_config)
self.load_lay_down_config()
# Store the node objects as node attributes
# (This is so we can access them as objects)
@@ -1027,7 +1032,7 @@ class Primaite(Env):
acl_rule_destination = item["destination"]
acl_rule_protocol = item["protocol"]
acl_rule_port = item["port"]
acl_rule_position = item["position"]
acl_rule_position = item.get("position")
self.acl.add_rule(
acl_rule_permission,

18
src/primaite/main.py
View File

@@ -14,18 +14,26 @@ def run(
training_config_path: Optional[Union[str, Path]] = "",
lay_down_config_path: Optional[Union[str, Path]] = "",
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Run the PrimAITE Session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
session = PrimaiteSession(training_config_path, lay_down_config_path, session_path)
session = PrimaiteSession(
training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config
)
session.setup()
session.learn()

22
src/primaite/primaite_session.py
View File

@@ -34,6 +34,8 @@ class PrimaiteSession:
training_config_path: Optional[Union[str, Path]] = "",
lay_down_config_path: Optional[Union[str, Path]] = "",
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
The PrimaiteSession constructor.
@@ -44,12 +46,18 @@ class PrimaiteSession:
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
self._agent_session: AgentSessionABC = None # noqa
self.session_path: Path = session_path # noqa
self.timestamp_str: str = None # noqa
self.learning_path: Path = None # noqa
self.evaluation_path: Path = None # noqa
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
# check if session path is provided
if session_path is not None:
@@ -67,12 +75,14 @@ class PrimaiteSession:
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path, str]] = training_config_path
self._training_config: Final[TrainingConfig] = training_config.load(self._training_config_path)
self._training_config: Final[TrainingConfig] = training_config.load(
self._training_config_path, legacy_training_config
)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path) # noqa
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config) # noqa
def setup(self) -> None:
"""Performs the session setup."""
@@ -139,7 +149,13 @@ class PrimaiteSession:
elif self._training_config.agent_framework == AgentFramework.SB3:
_LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.SB3}")
# Stable Baselines3 Agent
self._agent_session = SB3Agent(self._training_config_path, self._lay_down_config_path, self.session_path)
self._agent_session = SB3Agent(
self._training_config_path,
self._lay_down_config_path,
self.session_path,
self.legacy_training_config,
self.legacy_lay_down_config,
)
elif self._training_config.agent_framework == AgentFramework.RLLIB:
_LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.RLLIB}")

170
tests/config/legacy_conversion/legacy_config_1_DDOS_BASIC.yaml Normal file
View File

@@ -0,0 +1,170 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
- itemType: ACTIONS
type: NODE
- itemType: STEPS
steps: 128
- itemType: PORTS
portsList:
- port: '80'
- itemType: SERVICES
serviceList:
- name: TCP
- itemType: NODE
id: '1'
name: PC1
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.2
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '2'
name: SERVER
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.3
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '3'
name: PC2
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.4
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '4'
name: SWITCH1
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.5
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '5'
name: SWITCH2
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.6
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '6'
name: SWITCH3
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.7
softwareState: GOOD
fileSystemState: GOOD
- itemType: LINK
id: '7'
name: link1
bandwidth: 1000000000
source: '1'
destination: '4'
- itemType: LINK
id: '8'
name: link2
bandwidth: 1000000000
source: '4'
destination: '2'
- itemType: LINK
id: '9'
name: link3
bandwidth: 1000000000
source: '2'
destination: '5'
- itemType: LINK
id: '10'
name: link4
bandwidth: 1000000000
source: '2'
destination: '6'
- itemType: LINK
id: '11'
name: link5
bandwidth: 1000000000
source: '5'
destination: '3'
- itemType: LINK
id: '12'
name: link6
bandwidth: 1000000000
source: '6'
destination: '3'
- itemType: GREEN_IER
id: '13'
startStep: 1
endStep: 128
load: 100000
protocol: TCP
port: '80'
source: '3'
destination: '2'
missionCriticality: 5
- itemType: RED_POL
id: '14'
startStep: 50
endStep: 50
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
protocol: TCP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
id: '15'
startStep: 60
endStep: 100
load: 1000000
protocol: TCP
port: '80'
source: '1'
destination: '2'
missionCriticality: 0
- itemType: RED_POL
id: '16'
startStep: 80
endStep: 80
targetNodeId: '2'
initiator: IER
type: SERVICE
protocol: TCP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: ACL_RULE
id: '17'
permission: ALLOW
source: ANY
destination: ANY
protocol: ANY
port: ANY

362
tests/config/legacy_conversion/legacy_config_2_DDOS_BASIC.yaml Normal file
View File

@@ -0,0 +1,362 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
- itemType: ACTIONS
type: NODE
- itemType: STEPS
steps: 128
- itemType: PORTS
portsList:
- port: '80'
- itemType: SERVICES
serviceList:
- name: TCP
- itemType: NODE
id: '1'
name: PC1
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.11
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '2'
name: PC2
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.12
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '3'
name: PC3
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.13
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '4'
name: PC4
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.20.14
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '5'
name: SWITCH1
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.2
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '6'
name: IDS
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.4
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '7'
name: SWITCH2
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.3
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '8'
name: LOP1
baseType: SERVICE
nodeType: LOP
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.12
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '9'
name: SERVER1
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.14
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '10'
name: SERVER2
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.20.15
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
id: '11'
name: link1
bandwidth: 1000000000
source: '1'
destination: '5'
- itemType: LINK
id: '12'
name: link2
bandwidth: 1000000000
source: '2'
destination: '5'
- itemType: LINK
id: '13'
name: link3
bandwidth: 1000000000
source: '3'
destination: '5'
- itemType: LINK
id: '14'
name: link4
bandwidth: 1000000000
source: '4'
destination: '5'
- itemType: LINK
id: '15'
name: link5
bandwidth: 1000000000
source: '5'
destination: '6'
- itemType: LINK
id: '16'
name: link6
bandwidth: 1000000000
source: '5'
destination: '8'
- itemType: LINK
id: '17'
name: link7
bandwidth: 1000000000
source: '6'
destination: '7'
- itemType: LINK
id: '18'
name: link8
bandwidth: 1000000000
source: '8'
destination: '7'
- itemType: LINK
id: '19'
name: link9
bandwidth: 1000000000
source: '7'
destination: '9'
- itemType: LINK
id: '20'
name: link10
bandwidth: 1000000000
source: '7'
destination: '10'
- itemType: GREEN_IER
id: '21'
startStep: 1
endStep: 128
load: 100000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
id: '22'
startStep: 1
endStep: 128
load: 100000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
id: '23'
startStep: 1
endStep: 128
load: 100000
protocol: TCP
port: '80'
source: '9'
destination: '3'
missionCriticality: 5
- itemType: GREEN_IER
id: '24'
startStep: 1
endStep: 128
load: 100000
protocol: TCP
port: '80'
source: '4'
destination: '10'
missionCriticality: 2
- itemType: ACL_RULE
id: '25'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '26'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '27'
permission: ALLOW
source: 192.168.10.13
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '28'
permission: ALLOW
source: 192.168.20.14
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '29'
permission: ALLOW
source: 192.168.10.14
destination: 192.168.10.13
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '30'
permission: DENY
source: 192.168.10.11
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '31'
permission: DENY
source: 192.168.10.12
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '32'
permission: DENY
source: 192.168.10.13
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '33'
permission: DENY
source: 192.168.20.14
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: RED_POL
id: '34'
startStep: 20
endStep: 20
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
protocol: TCP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_POL
id: '35'
startStep: 20
endStep: 20
targetNodeId: '2'
initiator: DIRECT
type: SERVICE
protocol: TCP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
id: '36'
startStep: 30
endStep: 128
load: 440000000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 0
- itemType: RED_IER
id: '37'
startStep: 30
endStep: 128
load: 440000000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 0
- itemType: RED_POL
id: '38'
startStep: 30
endStep: 30
targetNodeId: '9'
initiator: IER
type: SERVICE
protocol: TCP
state: OVERWHELMED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA

166
tests/config/legacy_conversion/legacy_config_3_DOS_VERY_BASIC.yaml Normal file
View File

@@ -0,0 +1,166 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
- itemType: ACTIONS
type: NODE
- itemType: STEPS
steps: 256
- itemType: PORTS
portsList:
- port: '80'
- itemType: SERVICES
serviceList:
- name: TCP
- itemType: NODE
id: '1'
name: PC1
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.2
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '2'
name: PC2
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.3
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '3'
name: SWITCH1
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.1.1
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '4'
name: SERVER1
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.4
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
id: '5'
name: link1
bandwidth: 1000000000
source: '1'
destination: '3'
- itemType: LINK
id: '6'
name: link2
bandwidth: 1000000000
source: '2'
destination: '3'
- itemType: LINK
id: '7'
name: link3
bandwidth: 1000000000
source: '3'
destination: '4'
- itemType: GREEN_IER
id: '8'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '1'
destination: '4'
missionCriticality: 1
- itemType: GREEN_IER
id: '9'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '2'
destination: '4'
missionCriticality: 1
- itemType: GREEN_IER
id: '10'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '4'
destination: '2'
missionCriticality: 5
- itemType: ACL_RULE
id: '11'
permission: ALLOW
source: 192.168.1.2
destination: 192.168.1.4
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '12'
permission: ALLOW
source: 192.168.1.3
destination: 192.168.1.4
protocol: TCP
port: 80
- itemType: ACL_RULE
id: '13'
permission: ALLOW
source: 192.168.1.4
destination: 192.168.1.3
protocol: TCP
port: 80
- itemType: RED_POL
id: '14'
startStep: 20
endStep: 20
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
protocol: TCP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
id: '15'
startStep: 30
endStep: 256
load: 10000000
protocol: TCP
port: '80'
source: '1'
destination: '4'
missionCriticality: 0
- itemType: RED_POL
id: '16'
startStep: 40
endStep: 40
targetNodeId: '4'
initiator: IER
type: SERVICE
protocol: TCP
state: OVERWHELMED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA

534
tests/config/legacy_conversion/legacy_config_5_DATA_MANIPULATION.yaml Normal file
View File

@@ -0,0 +1,534 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
- itemType: ACTIONS
type: NODE
- itemType: STEPS
steps: 256
- itemType: PORTS
portsList:
- port: '80'
- port: '1433'
- port: '53'
- itemType: SERVICES
serviceList:
- name: TCP
- name: TCP_SQL
- name: UDP
- itemType: NODE
id: '1'
name: CLIENT_1
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.11
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
id: '2'
name: CLIENT_2
baseType: SERVICE
nodeType: COMPUTER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.10.12
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
id: '3'
name: SWITCH_1
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.10.1
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '4'
name: SECURITY_SUITE
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.10
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
id: '5'
name: MANAGEMENT_CONSOLE
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.1.12
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
id: '6'
name: SWITCH_2
baseType: ACTIVE
nodeType: SWITCH
priority: P2
hardwareState: 'ON'
ipAddress: 192.168.2.1
softwareState: GOOD
fileSystemState: GOOD
- itemType: NODE
id: '7'
name: WEB_SERVER
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.2.10
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: TCP_SQL
port: '1433'
state: GOOD
- itemType: NODE
id: '8'
name: DATABASE_SERVER
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.2.14
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: TCP_SQL
port: '1433'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
id: '9'
name: BACKUP_SERVER
baseType: SERVICE
nodeType: SERVER
priority: P5
hardwareState: 'ON'
ipAddress: 192.168.2.16
softwareState: GOOD
fileSystemState: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
id: '10'
name: LINK_1
bandwidth: 1000000000
source: '1'
destination: '3'
- itemType: LINK
id: '11'
name: LINK_2
bandwidth: 1000000000
source: '2'
destination: '3'
- itemType: LINK
id: '12'
name: LINK_3
bandwidth: 1000000000
source: '3'
destination: '4'
- itemType: LINK
id: '13'
name: LINK_4
bandwidth: 1000000000
source: '3'
destination: '5'
- itemType: LINK
id: '14'
name: LINK_5
bandwidth: 1000000000
source: '4'
destination: '6'
- itemType: LINK
id: '15'
name: LINK_6
bandwidth: 1000000000
source: '5'
destination: '6'
- itemType: LINK
id: '16'
name: LINK_7
bandwidth: 1000000000
source: '6'
destination: '7'
- itemType: LINK
id: '17'
name: LINK_8
bandwidth: 1000000000
source: '6'
destination: '8'
- itemType: LINK
id: '18'
name: LINK_9
bandwidth: 1000000000
source: '6'
destination: '9'
- itemType: GREEN_IER
id: '19'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '1'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
id: '20'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '7'
destination: '1'
missionCriticality: 5
- itemType: GREEN_IER
id: '21'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '2'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
id: '22'
startStep: 1
endStep: 256
load: 10000
protocol: TCP
port: '80'
source: '7'
destination: '2'
missionCriticality: 5
- itemType: GREEN_IER
id: '23'
startStep: 1
endStep: 256
load: 5000
protocol: TCP_SQL
port: '1433'
source: '7'
destination: '8'
missionCriticality: 5
- itemType: GREEN_IER
id: '24'
startStep: 1
endStep: 256
load: 100000
protocol: TCP_SQL
port: '1433'
source: '8'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
id: '25'
startStep: 1
endStep: 256
load: 50000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
id: '26'
startStep: 1
endStep: 256
load: 50000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
id: '27'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '7'
missionCriticality: 1
- itemType: GREEN_IER
id: '28'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '7'
destination: '5'
missionCriticality: 1
- itemType: GREEN_IER
id: '29'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '8'
missionCriticality: 1
- itemType: GREEN_IER
id: '30'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '8'
destination: '5'
missionCriticality: 1
- itemType: GREEN_IER
id: '31'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '9'
missionCriticality: 1
- itemType: GREEN_IER
id: '32'
startStep: 1
endStep: 256
load: 5000
protocol: TCP
port: '80'
source: '9'
destination: '5'
missionCriticality: 1
- itemType: ACL_RULE
id: '33'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '34'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '35'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '36'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '37'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.10.11
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '38'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.10.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '39'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '40'
permission: ALLOW
source: 192.168.2.14
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '41'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '42'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '43'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '44'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '45'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '46'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '47'
permission: ALLOW
source: 192.168.2.14
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '48'
permission: ALLOW
source: 192.168.2.16
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
id: '49'
permission: DENY
source: ANY
destination: ANY
protocol: ANY
port: ANY
- itemType: RED_POL
id: '50'
startStep: 50
endStep: 50
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
protocol: UDP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
id: '51'
startStep: 75
endStep: 105
load: 10000
protocol: UDP
port: '53'
source: '1'
destination: '8'
missionCriticality: 0
- itemType: RED_POL
id: '52'
startStep: 100
endStep: 100
targetNodeId: '8'
initiator: IER
type: SERVICE
protocol: UDP
state: COMPROMISED
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_POL
id: '53'
startStep: 105
endStep: 105
targetNodeId: '8'
initiator: SERVICE
type: FILE
protocol: NA
state: CORRUPT
sourceNodeId: '8'
sourceNodeService: UDP
sourceNodeServiceState: COMPROMISED
- itemType: RED_POL
id: '54'
startStep: 105
endStep: 105
targetNodeId: '8'
initiator: SERVICE
type: SERVICE
protocol: TCP_SQL
state: COMPROMISED
sourceNodeId: '8'
sourceNodeService: UDP
sourceNodeServiceState: COMPROMISED
- itemType: RED_POL
id: '55'
startStep: 125
endStep: 125
targetNodeId: '7'
initiator: SERVICE
type: SERVICE
protocol: TCP
state: OVERWHELMED
sourceNodeId: '8'
sourceNodeService: TCP_SQL
sourceNodeServiceState: COMPROMISED

8
tests/config/legacy_conversion/new_training_config.yaml
View File

@@ -21,12 +21,20 @@ agent_identifier: PPO
# "ACL"
# "ANY" node and acl actions
action_type: ANY
# Number of episodes for training to run per session
num_train_episodes: 10
# Number of time_steps for training per episode
num_train_steps: 256
# Number of episodes for evaluation to run per session
num_eval_episodes: 1
# Number of time_steps for evaluation per episode
num_eval_steps: 256
# Time delay between steps (for generic agents)
time_delay: 10
# Type of session to be run (TRAINING or EVALUATION)

29
tests/test_full_legacy_config_session.py Normal file
View File

@@ -0,0 +1,29 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import pytest
from primaite.main import run
from tests import TEST_CONFIG_ROOT
@pytest.mark.parametrize(
"legacy_file",
[
("legacy_config_1_DDOS_BASIC.yaml"),
("legacy_config_2_DDOS_BASIC.yaml"),
("legacy_config_3_DOS_VERY_BASIC.yaml"),
("legacy_config_5_DATA_MANIPULATION.yaml"),
],
)
def test_legacy_training_config_run_session(legacy_file):
"""Tests using legacy training and lay down config files in PrimAITE session end-to-end."""
legacy_training_config_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml"
legacy_lay_down_config_path = TEST_CONFIG_ROOT / "legacy_conversion" / legacy_file
# Run a PrimAITE session using legacy training and lay down config file paths
run(
legacy_training_config_path,
legacy_lay_down_config_path,
legacy_training_config=True,
legacy_lay_down_config=True,
)

44
tests/test_lay_down_config.py Normal file
View File

@@ -0,0 +1,44 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.lay_down_config import (
convert_legacy_lay_down_config,
data_manipulation_config_path,
ddos_basic_one_config_path,
ddos_basic_two_config_path,
dos_very_basic_config_path,
)
from tests import TEST_CONFIG_ROOT
@pytest.mark.parametrize(
"legacy_file, new_path",
[
("legacy_config_1_DDOS_BASIC.yaml", ddos_basic_one_config_path()),
("legacy_config_2_DDOS_BASIC.yaml", ddos_basic_two_config_path()),
("legacy_config_3_DOS_VERY_BASIC.yaml", dos_very_basic_config_path()),
("legacy_config_5_DATA_MANIPULATION.yaml", data_manipulation_config_path()),
],
)
def test_legacy_lay_down_config_load(legacy_file, new_path):
"""Tests converting legacy lay down files into the new format."""
with open(TEST_CONFIG_ROOT / "legacy_conversion" / legacy_file, "r") as file:
legacy_lay_down_config = yaml.safe_load(file)
with open(new_path, "r") as file:
new_lay_down_config = yaml.safe_load(file)
converted_lay_down_config = convert_legacy_lay_down_config(legacy_lay_down_config)
assert len(converted_lay_down_config) == len(new_lay_down_config)
for i, new_item in enumerate(new_lay_down_config):
converted_item = converted_lay_down_config[i]
for key, val in new_item.items():
if key == "position":
continue
assert key in converted_item
assert val == converted_item[key]