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

901 - merged dev into my branch

Browse Source
This commit is contained in:
SunilSamra
2023-06-13 08:54:33 +01:00
parent ed8b53f5ef dc0349c37b
commit ce5d38a1cd
75 changed files with 3640 additions and 2306 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/VERSION
View File

@@ -1 +1 @@
1.2.1
2.0.0dev0

125
src/primaite/__init__.py Normal file
View File

@@ -0,0 +1,125 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import logging
import logging.config
import sys
from logging import Logger, StreamHandler
from logging.handlers import RotatingFileHandler
from pathlib import Path
from typing import Final
import pkg_resources
import yaml
from platformdirs import PlatformDirs
_PLATFORM_DIRS: Final[PlatformDirs] = PlatformDirs(appname="primaite")
"""An instance of `PlatformDirs` set with appname='primaite'."""
def _get_primaite_config():
config_path = _PLATFORM_DIRS.user_config_path / "primaite_config.yaml"
if not config_path.exists():
config_path = Path(
pkg_resources.resource_filename(
"primaite", "setup/_package_data/primaite_config.yaml"
)
)
with open(config_path, "r") as file:
primaite_config = yaml.safe_load(file)
log_level_map = {
"NOTSET": logging.NOTSET,
"DEBUG": logging.DEBUG,
"INFO": logging.INFO,
"WARN": logging.WARN,
"ERROR": logging.ERROR,
"CRITICAL": logging.CRITICAL
}
primaite_config["log_level"] = log_level_map[primaite_config["log_level"]]
return primaite_config
_PRIMAITE_CONFIG = _get_primaite_config()
_USER_DIRS: Final[Path] = Path.home() / "primaite"
"""The users home space for PrimAITE which is located at: ~/primaite."""
NOTEBOOKS_DIR: Final[Path] = _USER_DIRS / "notebooks"
"""
The path to the users notebooks directory as an instance of `Path` or
`PosixPath`, depending on the OS.
Users notebooks are stored at: ``~/primaite/notebooks``.
"""
USERS_CONFIG_DIR: Final[Path] = _USER_DIRS / "config"
"""
The path to the users config directory as an instance of `Path` or
`PosixPath`, depending on the OS.
Users config files are stored at: ``~/primaite/config``.
"""
SESSIONS_DIR: Final[Path] = _USER_DIRS / "sessions"
"""
The path to the users PrimAITE Sessions directory as an instance of `Path` or
`PosixPath`, depending on the OS.
Users PrimAITE Sessions are stored at: ``~/primaite/sessions``.
"""
# region Setup Logging
def _log_dir() -> Path:
if sys.platform == "win32":
dir_path = _PLATFORM_DIRS.user_data_path / "logs"
else:
dir_path = _PLATFORM_DIRS.user_log_path
return dir_path
LOG_DIR: Final[Path] = _log_dir()
"""The path to the app log directory as an instance of `Path` or `PosixPath`, depending on the OS."""
LOG_DIR.mkdir(exist_ok=True, parents=True)
LOG_PATH: Final[Path] = LOG_DIR / "primaite.log"
"""The primaite.log file path as an instance of `Path` or `PosixPath`, depending on the OS."""
_STREAM_HANDLER: Final[StreamHandler] = StreamHandler()
_FILE_HANDLER: Final[RotatingFileHandler] = RotatingFileHandler(
filename=LOG_PATH,
maxBytes=10485760, # 10MB
backupCount=9, # Max 100MB of logs
encoding="utf8",
)
_STREAM_HANDLER.setLevel(_PRIMAITE_CONFIG["log_level"])
_FILE_HANDLER.setLevel(_PRIMAITE_CONFIG["log_level"])
_LOG_FORMAT_STR: Final[str] = _PRIMAITE_CONFIG["logger_format"]
_STREAM_HANDLER.setFormatter(logging.Formatter(_LOG_FORMAT_STR))
_FILE_HANDLER.setFormatter(logging.Formatter(_LOG_FORMAT_STR))
_LOGGER = logging.getLogger(__name__)
_LOGGER.addHandler(_STREAM_HANDLER)
_LOGGER.addHandler(_FILE_HANDLER)
def getLogger(name: str) -> Logger:
"""
Get a PrimAITE logger.
:param name: The logger name. Use ``__name__``.
:return: An instance of :py:class:`logging.Logger` with the PrimAITE
logging config.
"""
logger = logging.getLogger(name)
logger.setLevel(_PRIMAITE_CONFIG["log_level"])
return logger
# endregion
with open(Path(__file__).parent.resolve() / "VERSION", "r") as file:
__version__ = file.readline()

57
src/primaite/acl/access_control_list.py
View File

@@ -11,12 +11,21 @@ _LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
class AccessControlList:
"""Access Control List class."""
def __init__(self, implicit_permission):
def __init__(self, implicit_permission, max_acl_rules):
"""Init."""
# A list of ACL Rules
self.acl: List[ACLRule] = []
# Implicit ALLOW or DENY firewall spec
# Last rule in the ACL list
self.acl_implicit_rule = implicit_permission
self.max_acl_rules: int
# Create implicit rule based on input
if self.acl_implicit_rule == "DENY":
implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY")
else:
implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY")
# Maximum number of ACL Rules in ACL
self.max_acl_rules: int = max_acl_rules
# A list of ACL Rules
self.acl: List[ACLRule] = [implicit_rule]
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
"""
@@ -62,21 +71,15 @@ class AccessControlList:
Returns:
Indicates block if all conditions are satisfied.
"""
for rule_key, rule_value in self.acl.items():
if self.check_address_match(
rule_value, _source_ip_address, _dest_ip_address
):
for rule in self.acl:
if self.check_address_match(rule, _source_ip_address, _dest_ip_address):
if (
rule_value.get_protocol() == _protocol
or rule_value.get_protocol() == "ANY"
) and (
str(rule_value.get_port()) == str(_port)
or rule_value.get_port() == "ANY"
):
rule.get_protocol() == _protocol or rule.get_protocol() == "ANY"
) and (str(rule.get_port()) == str(_port) or rule.get_port() == "ANY"):
# There's a matching rule. Get the permission
if rule_value.get_permission() == "DENY":
if rule.get_permission() == "DENY":
return True
elif rule_value.get_permission() == "ALLOW":
elif rule.get_permission() == "ALLOW":
return False
# If there has been no rule to allow the IER through, it will return a blocked signal by default
@@ -92,18 +95,26 @@ class AccessControlList:
_dest_ip: the destination IP address
_protocol: the protocol
_port: the port
_position: position to insert ACL rule into ACL list
_position: position to insert ACL rule into ACL list (starting from index 1 and NOT 0)
"""
position_index = int(_position) - 1
new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
if _position < self.max_acl_rules - 1 and _position < 0:
if _position < len(self.acl):
self.acl.insert(_position, new_rule)
if len(self.acl) < self.max_acl_rules:
if len(self.acl) > position_index > -1:
try:
self.acl.insert(position_index, new_rule)
except Exception:
_LOGGER.info(
f"New Rule could NOT be added to list at position {position_index}."
)
else:
print("check logic on this")
_LOGGER.info(
f"Position {position_index} is an invalid index for list and/or overwrites implicit firewall rule"
)
else:
_LOGGER.info(
f"Position {_position} is an invalid index for list/overwriting implicit firewall rule"
f"The ACL list is FULL."
f"The list of ACLs has length {len(self.acl)} and it has a max capacity of {self.max_acl_rules}."
)
def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port):

171
src/primaite/cli.py Normal file
View File

@@ -0,0 +1,171 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Provides a CLI using Typer as an entry point."""
import logging
import os
import shutil
from pathlib import Path
from enum import Enum
from typing import Optional
import pkg_resources
import typer
import yaml
from platformdirs import PlatformDirs
from typing_extensions import Annotated
app = typer.Typer()
@app.command()
def build_dirs():
"""Build the PrimAITE app directories."""
from primaite.setup import setup_app_dirs
setup_app_dirs.run()
@app.command()
def reset_notebooks(overwrite: bool = True):
"""
Force a reset of the demo notebooks in the users notebooks directory.
:param overwrite: If True, will overwrite existing demo notebooks.
"""
from primaite.setup import reset_demo_notebooks
reset_demo_notebooks.run(overwrite)
@app.command()
def logs(last_n: Annotated[int, typer.Option("-n")]):
"""
Print the PrimAITE log file.
:param last_n: The number of lines to print. Default value is 10.
"""
import re
from primaite import LOG_PATH
if os.path.isfile(LOG_PATH):
with open(LOG_PATH) as file:
lines = file.readlines()
for line in lines[-last_n:]:
print(re.sub(r"\n*", "", line))
_LogLevel = Enum("LogLevel", {k: k for k in logging._levelToName.values()}) # noqa
@app.command()
def log_level(level: Annotated[Optional[_LogLevel], typer.Argument()] = None):
"""
View or set the PrimAITE Log Level.
To View, simply call: primaite log-level
To set, call: primaite log-level <desired log level>
For example, to set the to debug, call: primaite log-level DEBUG
"""
app_dirs = PlatformDirs(appname="primaite")
app_dirs.user_config_path.mkdir(exist_ok=True, parents=True)
user_config_path = app_dirs.user_config_path / "primaite_config.yaml"
if user_config_path.exists():
with open(user_config_path, "r") as file:
primaite_config = yaml.safe_load(file)
if level:
primaite_config["log_level"] = level.value
with open(user_config_path, "w") as file:
yaml.dump(primaite_config, file)
else:
level = primaite_config["log_level"]
print(f"PrimAITE Log Level: {level}")
@app.command()
def notebooks():
"""Start Jupyter Lab in the users PrimAITE notebooks directory."""
from primaite.notebooks import start_jupyter_session
start_jupyter_session()
@app.command()
def version():
"""Get the installed PrimAITE version number."""
import primaite
print(primaite.__version__)
@app.command()
def clean_up():
"""Cleans up left over files from previous version installations."""
from primaite.setup import old_installation_clean_up
old_installation_clean_up.run()
@app.command()
def setup(overwrite_existing: bool = True):
"""
Perform the PrimAITE first-time setup.
WARNING: All user-data will be lost.
"""
# Does this way to avoid using PrimAITE package before config is loaded
app_dirs = PlatformDirs(appname="primaite")
app_dirs.user_config_path.mkdir(exist_ok=True, parents=True)
user_config_path = app_dirs.user_config_path / "primaite_config.yaml"
build_config = overwrite_existing or (not user_config_path.exists())
if build_config:
pkg_config_path = Path(
pkg_resources.resource_filename(
"primaite", "setup/_package_data/primaite_config.yaml"
)
)
shutil.copy2(pkg_config_path, user_config_path)
from primaite import getLogger
from primaite.setup import (
old_installation_clean_up,
reset_demo_notebooks,
reset_example_configs,
setup_app_dirs,
)
_LOGGER = getLogger(__name__)
_LOGGER.info("Performing the PrimAITE first-time setup...")
if build_config:
_LOGGER.info("Building primaite_config.yaml...")
_LOGGER.info("Building the PrimAITE app directories...")
setup_app_dirs.run()
_LOGGER.info("Rebuilding the demo notebooks...")
reset_demo_notebooks.run(overwrite_existing=True)
_LOGGER.info("Rebuilding the example notebooks...")
reset_example_configs.run(overwrite_existing=True)
_LOGGER.info("Performing a clean-up of previous PrimAITE installations...")
old_installation_clean_up.run()
_LOGGER.info("PrimAITE setup complete!")
@app.command()
def session(tc: str, ldc: str):
"""
Run a PrimAITE session.
:param tc: The training config filepath.
:param ldc: The lay down config file path.
"""
from primaite.main import run
run(training_config_path=tc, lay_down_config_path=ldc)

91
src/primaite/common/config_values_main.py
View File

@@ -1,91 +0,0 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The config class."""
class ConfigValuesMain(object):
"""Class to hold main config values."""
def __init__(self):
"""Init."""
# Generic
self.agent_identifier = "" # the agent in use
self.observation_config = None # observation space config
self.num_episodes = 0 # number of episodes to train over
self.num_steps = 0 # number of steps in an episode
self.time_delay = 0 # delay between steps (ms) - applies to generic agents only
self.config_filename_use_case = "" # the filename for the Use Case config file
self.session_type = "" # the session type to run (TRAINING or EVALUATION)
# Environment
self.observation_space_high_value = (
0 # The high value for the observation space
)
# Reward values
# Generic
self.all_ok = 0
# Node Hardware State
self.off_should_be_on = 0
self.off_should_be_resetting = 0
self.on_should_be_off = 0
self.on_should_be_resetting = 0
self.resetting_should_be_on = 0
self.resetting_should_be_off = 0
self.resetting = 0
# Node Software or Service State
self.good_should_be_patching = 0
self.good_should_be_compromised = 0
self.good_should_be_overwhelmed = 0
self.patching_should_be_good = 0
self.patching_should_be_compromised = 0
self.patching_should_be_overwhelmed = 0
self.patching = 0
self.compromised_should_be_good = 0
self.compromised_should_be_patching = 0
self.compromised_should_be_overwhelmed = 0
self.compromised = 0
self.overwhelmed_should_be_good = 0
self.overwhelmed_should_be_patching = 0
self.overwhelmed_should_be_compromised = 0
self.overwhelmed = 0
# Node File System State
self.good_should_be_repairing = 0
self.good_should_be_restoring = 0
self.good_should_be_corrupt = 0
self.good_should_be_destroyed = 0
self.repairing_should_be_good = 0
self.repairing_should_be_restoring = 0
self.repairing_should_be_corrupt = 0
self.repairing_should_be_destroyed = (
0 # Repairing does not fix destroyed state - you need to restore
)
self.repairing = 0
self.restoring_should_be_good = 0
self.restoring_should_be_repairing = 0
self.restoring_should_be_corrupt = (
0 # Not the optimal method (as repair will fix corruption)
)
self.restoring_should_be_destroyed = 0
self.restoring = 0
self.corrupt_should_be_good = 0
self.corrupt_should_be_repairing = 0
self.corrupt_should_be_restoring = 0
self.corrupt_should_be_destroyed = 0
self.corrupt = 0
self.destroyed_should_be_good = 0
self.destroyed_should_be_repairing = 0
self.destroyed_should_be_restoring = 0
self.destroyed_should_be_corrupt = 0
self.destroyed = 0
self.scanning = 0
# IER status
self.red_ier_running = 0
self.green_ier_blocked = 0
# Patching / Reset
self.os_patching_duration = 0 # The time taken to patch the OS
self.node_reset_duration = 0 # The time taken to reset a node (hardware)
self.service_patching_duration = 0 # The time taken to patch a service
self.file_system_repairing_limit = 0 # The time take to repair a file
self.file_system_restoring_limit = 0 # The time take to restore a file
self.file_system_scanning_limit = 0 # The time taken to scan the file system

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

@@ -35,6 +35,8 @@ class HardwareState(Enum):
ON = 1
OFF = 2
RESETTING = 3
SHUTTING_DOWN = 4
BOOTING = 5
class SoftwareState(Enum):

95
src/primaite/common/training_config.py Normal file
View File

@@ -0,0 +1,95 @@
# # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
# """The config class."""
# from dataclasses import dataclass
#
# from primaite.common.enums import ActionType
#
#
# @dataclass()
# class TrainingConfig:
# """Class to hold main config values."""
#
# # Generic
# agent_identifier: str # The Red Agent algo/class to be used
# action_type: ActionType # type of action to use (NODE/ACL/ANY)
# num_episodes: int # number of episodes to train over
# num_steps: int # number of steps in an episode
# time_delay: int # delay between steps (ms) - applies to generic agents only
# # file
# session_type: str # the session type to run (TRAINING or EVALUATION)
# load_agent: str # Determine whether to load an agent from file
# agent_load_file: str # File path and file name of agent if you're loading one in
#
# # Environment
# observation_space_high_value: int # The high value for the observation space
#
# # Reward values
# # Generic
# all_ok: int
# # Node Hardware State
# off_should_be_on: int
# off_should_be_resetting: int
# on_should_be_off: int
# on_should_be_resetting: int
# resetting_should_be_on: int
# resetting_should_be_off: int
# resetting: int
# # Node Software or Service State
# good_should_be_patching: int
# good_should_be_compromised: int
# good_should_be_overwhelmed: int
# patching_should_be_good: int
# patching_should_be_compromised: int
# patching_should_be_overwhelmed: int
# patching: int
# compromised_should_be_good: int
# compromised_should_be_patching: int
# compromised_should_be_overwhelmed: int
# compromised: int
# overwhelmed_should_be_good: int
# overwhelmed_should_be_patching: int
# overwhelmed_should_be_compromised: int
# overwhelmed: int
# # Node File System State
# good_should_be_repairing: int
# good_should_be_restoring: int
# good_should_be_corrupt: int
# good_should_be_destroyed: int
# repairing_should_be_good: int
# repairing_should_be_restoring: int
# repairing_should_be_corrupt: int
# repairing_should_be_destroyed: int # Repairing does not fix destroyed state - you need to restore
#
# repairing: int
# restoring_should_be_good: int
# restoring_should_be_repairing: int
# restoring_should_be_corrupt: int # Not the optimal method (as repair will fix corruption)
#
# restoring_should_be_destroyed: int
# restoring: int
# corrupt_should_be_good: int
# corrupt_should_be_repairing: int
# corrupt_should_be_restoring: int
# corrupt_should_be_destroyed: int
# corrupt: int
# destroyed_should_be_good: int
# destroyed_should_be_repairing: int
# destroyed_should_be_restoring: int
# destroyed_should_be_corrupt: int
# destroyed: int
# scanning: int
# # IER status
# red_ier_running: int
# green_ier_blocked: int
#
# # Patching / Reset
# os_patching_duration: int # The time taken to patch the OS
# node_reset_duration: int # The time taken to reset a node (hardware)
# node_booting_duration = 0 # The Time taken to turn on the node
# node_shutdown_duration = 0 # The time taken to turn off the node
# service_patching_duration: int # The time taken to patch a service
# file_system_repairing_limit: int # The time take to repair a file
# file_system_restoring_limit: int # The time take to restore a file
# file_system_scanning_limit: int # The time taken to scan the file system
# # Patching / Reset
#

66
src/primaite/config/config_1_DDOS_BASIC.yaml → src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml
View File

@@ -1,14 +1,14 @@
- itemType: ACTIONS
- item_type: ACTIONS
type: NODE
- itemType: STEPS
- item_type: STEPS
steps: 128
- itemType: PORTS
portsList:
- item_type: PORTS
ports_list:
- port: '80'
- itemType: SERVICES
serviceList:
- item_type: SERVICES
service_list:
- name: TCP
- itemType: NODE
- item_type: NODE
node_id: '1'
name: PC1
node_class: SERVICE
@@ -22,7 +22,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '2'
name: SERVER
node_class: SERVICE
@@ -36,7 +36,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '3'
name: PC2
node_class: SERVICE
@@ -50,7 +50,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '4'
name: SWITCH1
node_class: ACTIVE
@@ -60,7 +60,7 @@
ip_address: 192.168.1.5
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '5'
name: SWITCH2
node_class: ACTIVE
@@ -70,7 +70,7 @@
ip_address: 192.168.1.6
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '6'
name: SWITCH3
node_class: ACTIVE
@@ -80,56 +80,56 @@
ip_address: 192.168.1.7
software_state: GOOD
file_system_state: GOOD
- itemType: LINK
- item_type: LINK
id: '7'
name: link1
bandwidth: 1000000000
source: '1'
destination: '4'
- itemType: LINK
- item_type: LINK
id: '8'
name: link2
bandwidth: 1000000000
source: '4'
destination: '2'
- itemType: LINK
- item_type: LINK
id: '9'
name: link3
bandwidth: 1000000000
source: '2'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '10'
name: link4
bandwidth: 1000000000
source: '2'
destination: '6'
- itemType: LINK
- item_type: LINK
id: '11'
name: link5
bandwidth: 1000000000
source: '5'
destination: '3'
- itemType: LINK
- item_type: LINK
id: '12'
name: link6
bandwidth: 1000000000
source: '6'
destination: '3'
- itemType: GREEN_IER
- item_type: GREEN_IER
id: '13'
startStep: 1
endStep: 128
start_step: 1
end_step: 128
load: 100000
protocol: TCP
port: '80'
source: '3'
destination: '2'
missionCriticality: 5
- itemType: RED_POL
mission_criticality: 5
- item_type: RED_POL
id: '14'
startStep: 50
endStep: 50
start_step: 50
end_step: 50
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
@@ -138,20 +138,20 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
- item_type: RED_IER
id: '15'
startStep: 60
endStep: 100
start_step: 60
end_step: 100
load: 1000000
protocol: TCP
port: '80'
source: '1'
destination: '2'
missionCriticality: 0
- itemType: RED_POL
mission_criticality: 0
- item_type: RED_POL
id: '16'
startStep: 80
endStep: 80
start_step: 80
end_step: 80
targetNodeId: '2'
initiator: IER
type: SERVICE
@@ -160,7 +160,7 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '17'
permission: ALLOW
source: ANY

136
src/primaite/config/config_2_DDOS_BASIC.yaml → src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml
View File

@@ -1,14 +1,14 @@
- itemType: ACTIONS
- item_type: ACTIONS
type: NODE
- itemType: STEPS
- item_type: STEPS
steps: 128
- itemType: PORTS
portsList:
- item_type: PORTS
ports_list:
- port: '80'
- itemType: SERVICES
serviceList:
- item_type: SERVICES
service_list:
- name: TCP
- itemType: NODE
- item_type: NODE
node_id: '1'
name: PC1
node_class: SERVICE
@@ -22,7 +22,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '2'
name: PC2
node_class: SERVICE
@@ -36,7 +36,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '3'
name: PC3
node_class: SERVICE
@@ -50,7 +50,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '4'
name: PC4
node_class: SERVICE
@@ -64,7 +64,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '5'
name: SWITCH1
node_class: ACTIVE
@@ -74,7 +74,7 @@
ip_address: 192.168.1.2
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '6'
name: IDS
node_class: SERVICE
@@ -88,7 +88,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '7'
name: SWITCH2
node_class: ACTIVE
@@ -98,7 +98,7 @@
ip_address: 192.168.1.3
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '8'
name: LOP1
node_class: SERVICE
@@ -112,7 +112,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '9'
name: SERVER1
node_class: SERVICE
@@ -126,7 +126,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '10'
name: SERVER2
node_class: SERVICE
@@ -140,173 +140,173 @@
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
- item_type: LINK
id: '11'
name: link1
bandwidth: 1000000000
source: '1'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '12'
name: link2
bandwidth: 1000000000
source: '2'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '13'
name: link3
bandwidth: 1000000000
source: '3'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '14'
name: link4
bandwidth: 1000000000
source: '4'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '15'
name: link5
bandwidth: 1000000000
source: '5'
destination: '6'
- itemType: LINK
- item_type: LINK
id: '16'
name: link6
bandwidth: 1000000000
source: '5'
destination: '8'
- itemType: LINK
- item_type: LINK
id: '17'
name: link7
bandwidth: 1000000000
source: '6'
destination: '7'
- itemType: LINK
- item_type: LINK
id: '18'
name: link8
bandwidth: 1000000000
source: '8'
destination: '7'
- itemType: LINK
- item_type: LINK
id: '19'
name: link9
bandwidth: 1000000000
source: '7'
destination: '9'
- itemType: LINK
- item_type: LINK
id: '20'
name: link10
bandwidth: 1000000000
source: '7'
destination: '10'
- itemType: GREEN_IER
- item_type: GREEN_IER
id: '21'
startStep: 1
endStep: 128
start_step: 1
end_step: 128
load: 100000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
mission_criticality: 2
- item_type: GREEN_IER
id: '22'
startStep: 1
endStep: 128
start_step: 1
end_step: 128
load: 100000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
mission_criticality: 2
- item_type: GREEN_IER
id: '23'
startStep: 1
endStep: 128
start_step: 1
end_step: 128
load: 100000
protocol: TCP
port: '80'
source: '9'
destination: '3'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '24'
startStep: 1
endStep: 128
start_step: 1
end_step: 128
load: 100000
protocol: TCP
port: '80'
source: '4'
destination: '10'
missionCriticality: 2
- itemType: ACL_RULE
mission_criticality: 2
- item_type: ACL_RULE
id: '25'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '26'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '27'
permission: ALLOW
source: 192.168.10.13
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '28'
permission: ALLOW
source: 192.168.20.14
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '29'
permission: ALLOW
source: 192.168.10.14
destination: 192.168.10.13
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '30'
permission: DENY
source: 192.168.10.11
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '31'
permission: DENY
source: 192.168.10.12
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '32'
permission: DENY
source: 192.168.10.13
destination: 192.168.20.15
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '33'
permission: DENY
source: 192.168.20.14
destination: 192.168.10.14
protocol: TCP
port: 80
- itemType: RED_POL
- item_type: RED_POL
id: '34'
startStep: 20
endStep: 20
start_step: 20
end_step: 20
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
@@ -315,10 +315,10 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_POL
- item_type: RED_POL
id: '35'
startStep: 20
endStep: 20
start_step: 20
end_step: 20
targetNodeId: '2'
initiator: DIRECT
type: SERVICE
@@ -327,30 +327,30 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
- item_type: RED_IER
id: '36'
startStep: 30
endStep: 128
start_step: 30
end_step: 128
load: 440000000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 0
- itemType: RED_IER
mission_criticality: 0
- item_type: RED_IER
id: '37'
startStep: 30
endStep: 128
start_step: 30
end_step: 128
load: 440000000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 0
- itemType: RED_POL
mission_criticality: 0
- item_type: RED_POL
id: '38'
startStep: 30
endStep: 30
start_step: 30
end_step: 30
targetNodeId: '9'
initiator: IER
type: SERVICE

76
src/primaite/config/config_3_DOS_VERY_BASIC.yaml → src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml
View File

@@ -1,14 +1,14 @@
- itemType: ACTIONS
- item_type: ACTIONS
type: NODE
- itemType: STEPS
- item_type: STEPS
steps: 256
- itemType: PORTS
portsList:
- item_type: PORTS
ports_list:
- port: '80'
- itemType: SERVICES
serviceList:
- item_type: SERVICES
service_list:
- name: TCP
- itemType: NODE
- item_type: NODE
node_id: '1'
name: PC1
node_class: SERVICE
@@ -22,7 +22,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '2'
name: PC2
node_class: SERVICE
@@ -36,7 +36,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '3'
name: SWITCH1
node_class: ACTIVE
@@ -46,7 +46,7 @@
ip_address: 192.168.1.1
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '4'
name: SERVER1
node_class: SERVICE
@@ -60,79 +60,79 @@
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
- item_type: LINK
id: '5'
name: link1
bandwidth: 1000000000
source: '1'
destination: '3'
- itemType: LINK
- item_type: LINK
id: '6'
name: link2
bandwidth: 1000000000
source: '2'
destination: '3'
- itemType: LINK
- item_type: LINK
id: '7'
name: link3
bandwidth: 1000000000
source: '3'
destination: '4'
- itemType: GREEN_IER
- item_type: GREEN_IER
id: '8'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '1'
destination: '4'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '9'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '2'
destination: '4'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '10'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '4'
destination: '2'
missionCriticality: 5
- itemType: ACL_RULE
mission_criticality: 5
- item_type: ACL_RULE
id: '11'
permission: ALLOW
source: 192.168.1.2
destination: 192.168.1.4
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '12'
permission: ALLOW
source: 192.168.1.3
destination: 192.168.1.4
protocol: TCP
port: 80
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '13'
permission: ALLOW
source: 192.168.1.4
destination: 192.168.1.3
protocol: TCP
port: 80
- itemType: RED_POL
- item_type: RED_POL
id: '14'
startStep: 20
endStep: 20
start_step: 20
end_step: 20
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
@@ -141,20 +141,20 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
- item_type: RED_IER
id: '15'
startStep: 30
endStep: 256
start_step: 30
end_step: 256
load: 10000000
protocol: TCP
port: '80'
source: '1'
destination: '4'
missionCriticality: 0
- itemType: RED_POL
mission_criticality: 0
- item_type: RED_POL
id: '16'
startStep: 40
endStep: 40
start_step: 40
end_step: 40
targetNodeId: '4'
initiator: IER
type: SERVICE

232
src/primaite/config/config_5_DATA_MANIPULATION.yaml → src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml
View File

@@ -1,18 +1,14 @@
- itemType: ACTIONS
type: NODE
- itemType: STEPS
steps: 256
- itemType: PORTS
portsList:
- item_type: PORTS
ports_list:
- port: '80'
- port: '1433'
- port: '53'
- itemType: SERVICES
serviceList:
- item_type: SERVICES
service_list:
- name: TCP
- name: TCP_SQL
- name: UDP
- itemType: NODE
- item_type: NODE
node_id: '1'
name: CLIENT_1
node_class: SERVICE
@@ -29,7 +25,7 @@
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '2'
name: CLIENT_2
node_class: SERVICE
@@ -43,7 +39,7 @@
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '3'
name: SWITCH_1
node_class: ACTIVE
@@ -53,7 +49,7 @@
ip_address: 192.168.10.1
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '4'
name: SECURITY_SUITE
node_class: SERVICE
@@ -70,7 +66,7 @@
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '5'
name: MANAGEMENT_CONSOLE
node_class: SERVICE
@@ -87,7 +83,7 @@
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '6'
name: SWITCH_2
node_class: ACTIVE
@@ -97,7 +93,7 @@
ip_address: 192.168.2.1
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '7'
name: WEB_SERVER
node_class: SERVICE
@@ -114,7 +110,7 @@
- name: TCP_SQL
port: '1433'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '8'
name: DATABASE_SERVER
node_class: SERVICE
@@ -134,7 +130,7 @@
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
- item_type: NODE
node_id: '9'
name: BACKUP_SERVER
node_class: SERVICE
@@ -148,323 +144,323 @@
- name: TCP
port: '80'
state: GOOD
- itemType: LINK
- item_type: LINK
id: '10'
name: LINK_1
bandwidth: 1000000000
source: '1'
destination: '3'
- itemType: LINK
- item_type: LINK
id: '11'
name: LINK_2
bandwidth: 1000000000
source: '2'
destination: '3'
- itemType: LINK
- item_type: LINK
id: '12'
name: LINK_3
bandwidth: 1000000000
source: '3'
destination: '4'
- itemType: LINK
- item_type: LINK
id: '13'
name: LINK_4
bandwidth: 1000000000
source: '3'
destination: '5'
- itemType: LINK
- item_type: LINK
id: '14'
name: LINK_5
bandwidth: 1000000000
source: '4'
destination: '6'
- itemType: LINK
- item_type: LINK
id: '15'
name: LINK_6
bandwidth: 1000000000
source: '5'
destination: '6'
- itemType: LINK
- item_type: LINK
id: '16'
name: LINK_7
bandwidth: 1000000000
source: '6'
destination: '7'
- itemType: LINK
- item_type: LINK
id: '17'
name: LINK_8
bandwidth: 1000000000
source: '6'
destination: '8'
- itemType: LINK
- item_type: LINK
id: '18'
name: LINK_9
bandwidth: 1000000000
source: '6'
destination: '9'
- itemType: GREEN_IER
- item_type: GREEN_IER
id: '19'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '1'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '20'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '7'
destination: '1'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '21'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '2'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '22'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 10000
protocol: TCP
port: '80'
source: '7'
destination: '2'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '23'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP_SQL
port: '1433'
source: '7'
destination: '8'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '24'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 100000
protocol: TCP_SQL
port: '1433'
source: '8'
destination: '7'
missionCriticality: 5
- itemType: GREEN_IER
mission_criticality: 5
- item_type: GREEN_IER
id: '25'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 50000
protocol: TCP
port: '80'
source: '1'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
mission_criticality: 2
- item_type: GREEN_IER
id: '26'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 50000
protocol: TCP
port: '80'
source: '2'
destination: '9'
missionCriticality: 2
- itemType: GREEN_IER
mission_criticality: 2
- item_type: GREEN_IER
id: '27'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '7'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '28'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '7'
destination: '5'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '29'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '8'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '30'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '8'
destination: '5'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '31'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '5'
destination: '9'
missionCriticality: 1
- itemType: GREEN_IER
mission_criticality: 1
- item_type: GREEN_IER
id: '32'
startStep: 1
endStep: 256
start_step: 1
end_step: 256
load: 5000
protocol: TCP
port: '80'
source: '9'
destination: '5'
missionCriticality: 1
- itemType: ACL_RULE
mission_criticality: 1
- item_type: ACL_RULE
id: '33'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '34'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '35'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '36'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '37'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.10.11
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '38'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.10.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '39'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '40'
permission: ALLOW
source: 192.168.2.14
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '41'
permission: ALLOW
source: 192.168.10.11
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '42'
permission: ALLOW
source: 192.168.10.12
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '43'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.10
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '44'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.14
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '45'
permission: ALLOW
source: 192.168.1.12
destination: 192.168.2.16
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '46'
permission: ALLOW
source: 192.168.2.10
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '47'
permission: ALLOW
source: 192.168.2.14
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '48'
permission: ALLOW
source: 192.168.2.16
destination: 192.168.1.12
protocol: ANY
port: ANY
- itemType: ACL_RULE
- item_type: ACL_RULE
id: '49'
permission: DENY
source: ANY
destination: ANY
protocol: ANY
port: ANY
- itemType: RED_POL
- item_type: RED_POL
id: '50'
startStep: 50
endStep: 50
start_step: 50
end_step: 50
targetNodeId: '1'
initiator: DIRECT
type: SERVICE
@@ -473,20 +469,20 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_IER
- item_type: RED_IER
id: '51'
startStep: 75
endStep: 105
start_step: 75
end_step: 105
load: 10000
protocol: UDP
port: '53'
source: '1'
destination: '8'
missionCriticality: 0
- itemType: RED_POL
mission_criticality: 0
- item_type: RED_POL
id: '52'
startStep: 100
endStep: 100
start_step: 100
end_step: 100
targetNodeId: '8'
initiator: IER
type: SERVICE
@@ -495,10 +491,10 @@
sourceNodeId: NA
sourceNodeService: NA
sourceNodeServiceState: NA
- itemType: RED_POL
- item_type: RED_POL
id: '53'
startStep: 105
endStep: 105
start_step: 105
end_step: 105
targetNodeId: '8'
initiator: SERVICE
type: FILE
@@ -507,10 +503,10 @@
sourceNodeId: '8'
sourceNodeService: UDP
sourceNodeServiceState: COMPROMISED
- itemType: RED_POL
- item_type: RED_POL
id: '54'
startStep: 105
endStep: 105
start_step: 105
end_step: 105
targetNodeId: '8'
initiator: SERVICE
type: SERVICE
@@ -519,10 +515,10 @@
sourceNodeId: '8'
sourceNodeService: UDP
sourceNodeServiceState: COMPROMISED
- itemType: RED_POL
- item_type: RED_POL
id: '55'
startStep: 125
endStep: 125
start_step: 125
end_step: 125
targetNodeId: '7'
initiator: SERVICE
type: SERVICE

94
src/primaite/config/_package_data/training/training_config_main.yaml Normal file
View File

@@ -0,0 +1,94 @@
# Main Config File
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: STABLE_BASELINES3_A2C
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
# "ANY" node and acl actions
action_type: NODE
# Number of episodes to run per session
num_episodes: 10
# Number of time_steps per episode
num_steps: 256
# Time delay between steps (for generic agents)
time_delay: 10
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in
agent_load_file: C:\[Path]\[agent_saved_filename.zip]
# Environment config values
# The high value for the observation space
observation_space_high_value: 1000000000
# Reward values
# Generic
all_ok: 0
# Node Hardware State
off_should_be_on: -10
off_should_be_resetting: -5
on_should_be_off: -2
on_should_be_resetting: -5
resetting_should_be_on: -5
resetting_should_be_off: -2
resetting: -3
# Node Software or Service State
good_should_be_patching: 2
good_should_be_compromised: 5
good_should_be_overwhelmed: 5
patching_should_be_good: -5
patching_should_be_compromised: 2
patching_should_be_overwhelmed: 2
patching: -3
compromised_should_be_good: -20
compromised_should_be_patching: -20
compromised_should_be_overwhelmed: -20
compromised: -20
overwhelmed_should_be_good: -20
overwhelmed_should_be_patching: -20
overwhelmed_should_be_compromised: -20
overwhelmed: -20
# Node File System State
good_should_be_repairing: 2
good_should_be_restoring: 2
good_should_be_corrupt: 5
good_should_be_destroyed: 10
repairing_should_be_good: -5
repairing_should_be_restoring: 2
repairing_should_be_corrupt: 2
repairing_should_be_destroyed: 0
repairing: -3
restoring_should_be_good: -10
restoring_should_be_repairing: -2
restoring_should_be_corrupt: 1
restoring_should_be_destroyed: 2
restoring: -6
corrupt_should_be_good: -10
corrupt_should_be_repairing: -10
corrupt_should_be_restoring: -10
corrupt_should_be_destroyed: 2
corrupt: -10
destroyed_should_be_good: -20
destroyed_should_be_repairing: -20
destroyed_should_be_restoring: -20
destroyed_should_be_corrupt: -20
destroyed: -20
scanning: -2
# IER status
red_ier_running: -5
green_ier_blocked: -10
# Patching / Reset durations
os_patching_duration: 5 # The time taken to patch the OS
node_reset_duration: 5 # The time taken to reset a node (hardware)
service_patching_duration: 5 # The time taken to patch a service
file_system_repairing_limit: 5 # The time take to repair the file system
file_system_restoring_limit: 5 # The time take to restore the file system
file_system_scanning_limit: 5 # The time taken to scan the file system

533
src/primaite/config/config_UNIT_TEST.yaml
View File

@@ -1,533 +0,0 @@
- 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
node_id: '1'
name: CLIENT_1
node_class: SERVICE
node_type: COMPUTER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.10.11
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
node_id: '2'
name: CLIENT_2
node_class: SERVICE
node_type: COMPUTER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.10.12
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- itemType: NODE
node_id: '3'
name: SWITCH_1
node_class: ACTIVE
node_type: SWITCH
priority: P2
hardware_state: 'ON'
ip_address: 192.168.10.1
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
node_id: '4'
name: SECURITY_SUITE
node_class: SERVICE
node_type: SERVER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.1.10
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
node_id: '5'
name: MANAGEMENT_CONSOLE
node_class: SERVICE
node_type: SERVER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.1.12
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
node_id: '6'
name: SWITCH_2
node_class: ACTIVE
node_type: SWITCH
priority: P2
hardware_state: 'ON'
ip_address: 192.168.2.1
software_state: GOOD
file_system_state: GOOD
- itemType: NODE
node_id: '7'
name: WEB_SERVER
node_class: SERVICE
node_type: SERVER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.2.10
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: TCP_SQL
port: '1433'
state: GOOD
- itemType: NODE
node_id: '8'
name: DATABASE_SERVER
node_class: SERVICE
node_type: SERVER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.2.14
software_state: GOOD
file_system_state: GOOD
services:
- name: TCP
port: '80'
state: GOOD
- name: TCP_SQL
port: '1433'
state: GOOD
- name: UDP
port: '53'
state: GOOD
- itemType: NODE
node_id: '9'
name: BACKUP_SERVER
node_class: SERVICE
node_type: SERVER
priority: P5
hardware_state: 'ON'
ip_address: 192.168.2.16
software_state: GOOD
file_system_state: 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

89
src/primaite/config/config_main.yaml
View File

@@ -1,89 +0,0 @@
# Main Config File
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agentIdentifier: STABLE_BASELINES3_A2C
# Number of episodes to run per session
numEpisodes: 10
# Time delay between steps (for generic agents)
timeDelay: 10
# Filename of the scenario / laydown
configFilename: config_5_DATA_MANIPULATION.yaml
# Type of session to be run (TRAINING or EVALUATION)
sessionType: TRAINING
# Determine whether to load an agent from file
loadAgent: False
# File path and file name of agent if you're loading one in
agentLoadFile: C:\[Path]\[agent_saved_filename.zip]
# Environment config values
# The high value for the observation space
observationSpaceHighValue: 1000000000
# Reward values
# Generic
allOk: 0
# Node Hardware State
offShouldBeOn: -10
offShouldBeResetting: -5
onShouldBeOff: -2
onShouldBeResetting: -5
resettingShouldBeOn: -5
resettingShouldBeOff: -2
resetting: -3
# Node Software or Service State
goodShouldBePatching: 2
goodShouldBeCompromised: 5
goodShouldBeOverwhelmed: 5
patchingShouldBeGood: -5
patchingShouldBeCompromised: 2
patchingShouldBeOverwhelmed: 2
patching: -3
compromisedShouldBeGood: -20
compromisedShouldBePatching: -20
compromisedShouldBeOverwhelmed: -20
compromised: -20
overwhelmedShouldBeGood: -20
overwhelmedShouldBePatching: -20
overwhelmedShouldBeCompromised: -20
overwhelmed: -20
# Node File System State
goodShouldBeRepairing: 2
goodShouldBeRestoring: 2
goodShouldBeCorrupt: 5
goodShouldBeDestroyed: 10
repairingShouldBeGood: -5
repairingShouldBeRestoring: 2
repairingShouldBeCorrupt: 2
repairingShouldBeDestroyed: 0
repairing: -3
restoringShouldBeGood: -10
restoringShouldBeRepairing: -2
restoringShouldBeCorrupt: 1
restoringShouldBeDestroyed: 2
restoring: -6
corruptShouldBeGood: -10
corruptShouldBeRepairing: -10
corruptShouldBeRestoring: -10
corruptShouldBeDestroyed: 2
corrupt: -10
destroyedShouldBeGood: -20
destroyedShouldBeRepairing: -20
destroyedShouldBeRestoring: -20
destroyedShouldBeCorrupt: -20
destroyed: -20
scanning: -2
# IER status
redIerRunning: -5
greenIerBlocked: -10
# Patching / Reset durations
osPatchingDuration: 5 # The time taken to patch the OS
nodeResetDuration: 5 # The time taken to reset a node (hardware)
servicePatchingDuration: 5 # The time taken to patch a service
fileSystemRepairingLimit: 5 # The time take to repair the file system
fileSystemRestoringLimit: 5 # The time take to restore the file system
fileSystemScanningLimit: 5 # The time taken to scan the file system

69
src/primaite/config/lay_down_config.py Normal file
View File

@@ -0,0 +1,69 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from pathlib import Path
from typing import Final
from primaite import USERS_CONFIG_DIR, getLogger
_LOGGER = getLogger(__name__)
_EXAMPLE_LAY_DOWN: Final[Path] = USERS_CONFIG_DIR / "example_config" / "lay_down"
def ddos_basic_one_config_path() -> Path:
"""
The path to the example lay_down_config_1_DDOS_basic.yaml file.
:return: The file path.
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_1_DDOS_basic.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path
def ddos_basic_two_config_path() -> Path:
"""
The path to the example lay_down_config_2_DDOS_basic.yaml file.
:return: The file path.
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_2_DDOS_basic.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path
def dos_very_basic_config_path() -> Path:
"""
The path to the example lay_down_config_3_DOS_very_basic.yaml file.
:return: The file path.
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_3_DOS_very_basic.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path
def data_manipulation_config_path() -> Path:
"""
The path to the example lay_down_config_5_data_manipulation.yaml file.
:return: The file path.
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_5_data_manipulation.yaml"
if not path.exists():
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path

316
src/primaite/config/training_config.py Normal file
View File

@@ -0,0 +1,316 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Final, Union, Optional
import yaml
from primaite import USERS_CONFIG_DIR, getLogger
from primaite.common.enums import ActionType
_LOGGER = getLogger(__name__)
_EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training"
@dataclass()
class TrainingConfig:
"""The Training Config class."""
# Generic
agent_identifier: str = "STABLE_BASELINES3_A2C"
"The Red Agent algo/class to be used."
action_type: ActionType = ActionType.ANY
"The ActionType to use."
num_episodes: int = 10
"The number of episodes to train over."
num_steps: int = 256
"The number of steps in an episode."
observation_space: dict = field(
default_factory=lambda: {"components": [{"name": "NODE_LINK_TABLE"}]}
)
"The observation space config dict."
time_delay: int = 10
"The delay between steps (ms). Applies to generic agents only."
# file
session_type: str = "TRAINING"
"the session type to run (TRAINING or EVALUATION)"
load_agent: str = False
"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."
# Environment
observation_space_high_value: int = 1000000000
"The high value for the observation space."
# Reward values
# Generic
all_ok: int = 0
# Node Hardware State
off_should_be_on: int = -10
off_should_be_resetting: int = -5
on_should_be_off: int = -2
on_should_be_resetting: int = -5
resetting_should_be_on: int = -5
resetting_should_be_off: int = -2
resetting: int = -3
# Node Software or Service State
good_should_be_patching: int = 2
good_should_be_compromised: int = 5
good_should_be_overwhelmed: int = 5
patching_should_be_good: int = -5
patching_should_be_compromised: int = 2
patching_should_be_overwhelmed: int = 2
patching: int = -3
compromised_should_be_good: int = -20
compromised_should_be_patching: int = -20
compromised_should_be_overwhelmed: int = -20
compromised: int = -20
overwhelmed_should_be_good: int = -20
overwhelmed_should_be_patching: int = -20
overwhelmed_should_be_compromised: int = -20
overwhelmed: int = -20
# Node File System State
good_should_be_repairing: int = 2
good_should_be_restoring: int = 2
good_should_be_corrupt: int = 5
good_should_be_destroyed: int = 10
repairing_should_be_good: int = -5
repairing_should_be_restoring: int = 2
repairing_should_be_corrupt: int = 2
repairing_should_be_destroyed: int = 0
repairing: int = -3
restoring_should_be_good: int = -10
restoring_should_be_repairing: int = -2
restoring_should_be_corrupt: int = 1
restoring_should_be_destroyed: int = 2
restoring: int = -6
corrupt_should_be_good: int = -10
corrupt_should_be_repairing: int = -10
corrupt_should_be_restoring: int = -10
corrupt_should_be_destroyed: int = 2
corrupt: int = -10
destroyed_should_be_good: int = -20
destroyed_should_be_repairing: int = -20
destroyed_should_be_restoring: int = -20
destroyed_should_be_corrupt: int = -20
destroyed: int = -20
scanning: int = -2
# IER status
red_ier_running: int = -5
green_ier_blocked: int = -10
# Patching / Reset durations
os_patching_duration: int = 5
"The time taken to patch the OS."
node_reset_duration: int = 5
"The time taken to reset a node (hardware)."
node_booting_duration: int = 3
"The Time taken to turn on the node."
node_shutdown_duration: int = 2
"The time taken to turn off the node."
service_patching_duration: int = 5
"The time taken to patch a service."
file_system_repairing_limit: int = 5
"The time take to repair the file system."
file_system_restoring_limit: int = 5
"The time take to restore the file system."
file_system_scanning_limit: int = 5
"The time taken to scan the file system."
def to_dict(self, json_serializable: bool = True):
"""
Serialise the ``TrainingConfig`` as dict.
:param json_serializable: If True, Enums are converted to their
string name.
:return: The ``TrainingConfig`` as a dict.
"""
data = self.__dict__
if json_serializable:
data["action_type"] = self.action_type.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:
"""
Read in a training config yaml file.
:param file_path: The config file path.
:param legacy_file: True if the config file is legacy format, otherwise
False.
:return: An instance of
:class:`~primaite.config.training_config.TrainingConfig`.
:raises ValueError: If the file_path does not exist.
:raises TypeError: When the TrainingConfig object cannot be created
using the values from the config file read from ``file_path``.
"""
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 training config file: {file_path}")
if legacy_file:
try:
config = convert_legacy_training_config_dict(config)
except KeyError:
msg = (
f"Failed to convert training config file {file_path} "
f"from legacy format. Attempting to use file as is."
)
_LOGGER.error(msg)
# Convert values to Enums
config["action_type"] = ActionType[config["action_type"]]
try:
return TrainingConfig(**config)
except TypeError as e:
msg = (
f"Error when creating an instance of {TrainingConfig} "
f"from the training config file {file_path}"
)
_LOGGER.critical(msg, exc_info=True)
raise e
msg = f"Cannot load the training config as it does not exist: {file_path}"
_LOGGER.error(msg)
raise ValueError(msg)
def convert_legacy_training_config_dict(
legacy_config_dict: Dict[str, Any],
num_steps: int = 256,
action_type: str = "ANY"
) -> Dict[str, Any]:
"""
Convert a legacy training config dict to the new format.
:param legacy_config_dict: A legacy training config dict.
:param num_steps: The number of steps to set as legacy training configs
don't have num_steps values.
:param action_type: The action space type to set as legacy training configs
don't have action_type values.
:return: The converted training config dict.
"""
config_dict = {
"num_steps": num_steps,
"action_type": action_type
}
for legacy_key, value in legacy_config_dict.items():
new_key = _get_new_key_from_legacy(legacy_key)
if new_key:
config_dict[new_key] = value
return config_dict
def _get_new_key_from_legacy(legacy_key: str) -> str:
"""
Maps legacy training config keys to the new format keys.
:param legacy_key: A legacy training config key.
:return: The mapped key.
"""
key_mapping = {
"agentIdentifier": "agent_identifier",
"numEpisodes": "num_episodes",
"timeDelay": "time_delay",
"configFilename": None,
"sessionType": "session_type",
"loadAgent": "load_agent",
"agentLoadFile": "agent_load_file",
"observationSpaceHighValue": "observation_space_high_value",
"allOk": "all_ok",
"offShouldBeOn": "off_should_be_on",
"offShouldBeResetting": "off_should_be_resetting",
"onShouldBeOff": "on_should_be_off",
"onShouldBeResetting": "on_should_be_resetting",
"resettingShouldBeOn": "resetting_should_be_on",
"resettingShouldBeOff": "resetting_should_be_off",
"resetting": "resetting",
"goodShouldBePatching": "good_should_be_patching",
"goodShouldBeCompromised": "good_should_be_compromised",
"goodShouldBeOverwhelmed": "good_should_be_overwhelmed",
"patchingShouldBeGood": "patching_should_be_good",
"patchingShouldBeCompromised": "patching_should_be_compromised",
"patchingShouldBeOverwhelmed": "patching_should_be_overwhelmed",
"patching": "patching",
"compromisedShouldBeGood": "compromised_should_be_good",
"compromisedShouldBePatching": "compromised_should_be_patching",
"compromisedShouldBeOverwhelmed": "compromised_should_be_overwhelmed",
"compromised": "compromised",
"overwhelmedShouldBeGood": "overwhelmed_should_be_good",
"overwhelmedShouldBePatching": "overwhelmed_should_be_patching",
"overwhelmedShouldBeCompromised": "overwhelmed_should_be_compromised",
"overwhelmed": "overwhelmed",
"goodShouldBeRepairing": "good_should_be_repairing",
"goodShouldBeRestoring": "good_should_be_restoring",
"goodShouldBeCorrupt": "good_should_be_corrupt",
"goodShouldBeDestroyed": "good_should_be_destroyed",
"repairingShouldBeGood": "repairing_should_be_good",
"repairingShouldBeRestoring": "repairing_should_be_restoring",
"repairingShouldBeCorrupt": "repairing_should_be_corrupt",
"repairingShouldBeDestroyed": "repairing_should_be_destroyed",
"repairing": "repairing",
"restoringShouldBeGood": "restoring_should_be_good",
"restoringShouldBeRepairing": "restoring_should_be_repairing",
"restoringShouldBeCorrupt": "restoring_should_be_corrupt",
"restoringShouldBeDestroyed": "restoring_should_be_destroyed",
"restoring": "restoring",
"corruptShouldBeGood": "corrupt_should_be_good",
"corruptShouldBeRepairing": "corrupt_should_be_repairing",
"corruptShouldBeRestoring": "corrupt_should_be_restoring",
"corruptShouldBeDestroyed": "corrupt_should_be_destroyed",
"corrupt": "corrupt",
"destroyedShouldBeGood": "destroyed_should_be_good",
"destroyedShouldBeRepairing": "destroyed_should_be_repairing",
"destroyedShouldBeRestoring": "destroyed_should_be_restoring",
"destroyedShouldBeCorrupt": "destroyed_should_be_corrupt",
"destroyed": "destroyed",
"scanning": "scanning",
"redIerRunning": "red_ier_running",
"greenIerBlocked": "green_ier_blocked",
"osPatchingDuration": "os_patching_duration",
"nodeResetDuration": "node_reset_duration",
"nodeBootingDuration": "node_booting_duration",
"nodeShutdownDuration": "node_shutdown_duration",
"servicePatchingDuration": "service_patching_duration",
"fileSystemRepairingLimit": "file_system_repairing_limit",
"fileSystemRestoringLimit": "file_system_restoring_limit",
"fileSystemScanningLimit": "file_system_scanning_limit",
}
return key_mapping[legacy_key]

35
src/primaite/environment/observations.py
View File

@@ -355,21 +355,50 @@ class AccessControlList(AbstractObservationComponent):
# 3. Initialise observation with zeroes
self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
# Dictionary to map services to numbers for obs space
self.services_dict = {}
def update(self):
"""Update the observation based on current environment state.
The structure of the observation space is described in :class:`.AccessControlList`
"""
obs = []
for acl_rule in self.env.acl:
for acl_rule in self.env.acl.acl:
permission = acl_rule.permission
source_ip = acl_rule.source_ip
dest_ip = acl_rule.dest_ip
protocol = acl_rule.protocol
port = acl_rule.port
position = self.env.acl.index(acl_rule)
position = self.env.acl.acl.index(acl_rule)
if permission == "DENY":
permission_int = 0
else:
permission_int = 1
obs.extend([permission, source_ip, dest_ip, protocol, port, position])
if source_ip == "ANY":
source_ip = 0
if dest_ip == "ANY":
dest_ip = 0
if port == "ANY":
port = 0
if protocol == "ANY":
protocol_int = 0
else:
while True:
if protocol in self.service_dict:
protocol_int = self.services_dict[protocol]
break
else:
self.services_dict[protocol] = len(self.services_dict) + 1
continue
# [0 - DENY, 1 - ALLOW] Permission
# [0 - ANY, x - IP Address/Protocol/Port]
print(permission_int, source_ip, dest_ip, protocol_int, port)
obs.extend(
[permission_int, source_ip, dest_ip, protocol_int, port, position]
)
self.current_observation[:] = obs

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

@@ -1,12 +1,11 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Main environment module containing the PRIMmary AI Training Evironment (Primaite) class."""
import copy
import csv
import logging
import os.path
from datetime import datetime
from typing import Dict, Tuple
from pathlib import Path
from typing import Dict, Tuple, Union
import networkx as nx
import numpy as np
@@ -24,11 +23,13 @@ from primaite.common.enums import (
NodePOLInitiator,
NodePOLType,
NodeType,
ObservationType,
Priority,
RulePermissionType,
SoftwareState,
)
from primaite.common.service import Service
from primaite.config import training_config
from primaite.config.training_config import TrainingConfig
from primaite.environment.observations import ObservationsHandler
from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link
@@ -56,29 +57,41 @@ class Primaite(Env):
ACTION_SPACE_ACL_ACTION_VALUES: int = 3
ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2
def __init__(self, _config_values, _transaction_list):
def __init__(
self,
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
transaction_list,
session_path: Path,
timestamp_str: str,
):
"""
Init.
The Primaite constructor.
Args:
_episode_steps: The number of steps for the episode
_config_filename: The name of config file
_transaction_list: The list of transactions to populate
_agent_identifier: Identifier for the agent
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.
:param transaction_list: The list of transactions to populate.
: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>.
"""
super(Primaite, self).__init__()
self._training_config_path = training_config_path
self._lay_down_config_path = lay_down_config_path
# Take a copy of the config values
self.config_values = _config_values
self.training_config: TrainingConfig = training_config.load(
training_config_path
)
# Number of steps in an episode
self.episode_steps = 0
self.episode_steps = self.training_config.num_steps
super(Primaite, self).__init__()
# Transaction list
self.transaction_list = _transaction_list
self.transaction_list = transaction_list
# The agent in use
self.agent_identifier = self.config_values.agent_identifier
self.agent_identifier = self.training_config.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes: Dict[str, NodeUnion] = {}
@@ -122,6 +135,9 @@ class Primaite(Env):
# Create step count
self.step_count = 0
self.total_step_count: int = 0
"""The total number of time steps completed."""
# Create step info dictionary
self.step_info = {}
@@ -149,29 +165,21 @@ class Primaite(Env):
# The action type
self.action_type = 0
# TODO fix up with TrainingConfig
# stores the observation config from the yaml, default is NODE_LINK_TABLE
self.obs_config: dict = {"components": [{"name": "NODE_LINK_TABLE"}]}
if self.config_values.observation_config is not None:
self.obs_config = self.config_values.observation_config
if self.training_config.observation_space is not None:
self.obs_config = self.training_config.observation_space
# Observation Handler manages the user-configurable observation space.
# It will be initialised later.
self.obs_handler: ObservationsHandler
# Set by main_config
# Adds a DENY ALL or ALLOW ALL to the end of the Access Control List
self.acl_implicit_rule = RulePermissionType.DENY
# Sets a limit to how many ACL
self.max_acl_rules = 0
# Open the config file and build the environment laydown
try:
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:
_LOGGER.error("Could not load the environment configuration")
_LOGGER.error("Exception occured", exc_info=True)
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()
# Store the node objects as node attributes
# (This is so we can access them as objects)
@@ -188,15 +196,10 @@ class Primaite(Env):
try:
plt.tight_layout()
nx.draw_networkx(self.network, with_labels=True)
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
datetime.now() # current date and time
path = "outputs/diagrams"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/diagrams/network_" + time + ".png"
plt.savefig(filename, format="PNG")
file_path = session_path / f"network_{timestamp_str}.png"
plt.savefig(file_path, format="PNG")
plt.clf()
except Exception:
_LOGGER.error("Could not save network diagram")
@@ -207,7 +210,7 @@ class Primaite(Env):
self.observation_space, self.env_obs = self.init_observations()
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ActionType.NODE:
if self.training_config.action_type == ActionType.NODE:
_LOGGER.info("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
@@ -216,7 +219,7 @@ class Primaite(Env):
# [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa
self.action_dict = self.create_node_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.action_type == ActionType.ACL:
elif self.training_config.action_type == ActionType.ACL:
_LOGGER.info("Action space type ACL selected")
# Terms (for ACL action space):
# [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
@@ -227,25 +230,21 @@ class Primaite(Env):
# [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
self.action_dict = self.create_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.action_type == ActionType.ANY:
elif self.training_config.action_type == ActionType.ANY:
_LOGGER.info("Action space type ANY selected - Node + ACL")
self.action_dict = self.create_node_and_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
else:
_LOGGER.info("Invalid action type selected")
_LOGGER.info(
f"Invalid action type selected: {self.training_config.action_type}"
)
# Set up a csv to store the results of the training
try:
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
header = ["Episode", "Average Reward"]
# Check whether the output/rerults folder exists (doesn't exist by default install)
path = "outputs/results/"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/results/average_reward_per_episode_" + time + ".csv"
self.csv_file = open(filename, "w", encoding="UTF8", newline="")
file_name = f"average_reward_per_episode_{timestamp_str}.csv"
file_path = session_path / file_name
self.csv_file = open(file_path, "w", encoding="UTF8", newline="")
self.csv_writer = csv.writer(self.csv_file)
self.csv_writer.writerow(header)
except Exception:
@@ -304,7 +303,8 @@ class Primaite(Env):
done = False
self.step_count += 1
# print("Episode step: " + str(self.stepCount))
self.total_step_count += 1
# print("Episode step: " + str(self.step_count))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
@@ -376,13 +376,13 @@ class Primaite(Env):
self.green_iers,
self.red_iers,
self.step_count,
self.config_values,
self.training_config,
)
print(f" Step {self.step_count} 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
if self.config_values.session_type == "EVALUATION":
if self.training_config.session_type == "EVALUATION":
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
@@ -407,7 +407,6 @@ class Primaite(Env):
def __close__(self):
"""Override close function."""
self.csv_file.close()
self.config_file.close()
def init_acl(self):
"""Initialise the Access Control List."""
@@ -433,9 +432,10 @@ class Primaite(Env):
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.action_type == ActionType.NODE:
if self.training_config.action_type == ActionType.NODE:
self.apply_actions_to_nodes(_action)
elif self.action_type == ActionType.ACL:
elif self.training_config.action_type == ActionType.ACL:
self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 6
@@ -673,44 +673,39 @@ class Primaite(Env):
self.obs_handler.update_obs()
self.env_obs = self.obs_handler.current_observation
def load_config(self):
def load_lay_down_config(self):
"""Loads config data in order to build the environment configuration."""
for item in self.config_data:
if item["itemType"] == "NODE":
for item in self.lay_down_config:
if item["item_type"] == "NODE":
# Create a node
self.create_node(item)
elif item["itemType"] == "LINK":
elif item["item_type"] == "LINK":
# Create a link
self.create_link(item)
elif item["itemType"] == "GREEN_IER":
elif item["item_type"] == "GREEN_IER":
# Create a Green IER
self.create_green_ier(item)
elif item["itemType"] == "GREEN_POL":
elif item["item_type"] == "GREEN_POL":
# Create a Green PoL
self.create_green_pol(item)
elif item["itemType"] == "RED_IER":
elif item["item_type"] == "RED_IER":
# Create a Red IER
self.create_red_ier(item)
elif item["itemType"] == "RED_POL":
elif item["item_type"] == "RED_POL":
# Create a Red PoL
self.create_red_pol(item)
elif item["itemType"] == "ACL_RULE":
elif item["item_type"] == "ACL_RULE":
# Create an ACL rule
self.create_acl_rule(item)
elif item["itemType"] == "SERVICES":
elif item["item_type"] == "SERVICES":
# Create the list of services
self.create_services_list(item)
elif item["itemType"] == "PORTS":
elif item["item_type"] == "PORTS":
# Create the list of ports
self.create_ports_list(item)
elif item["itemType"] == "ACTIONS":
# Get the action information
self.get_action_info(item)
elif item["itemType"] == "STEPS":
# Get the steps information
self.get_steps_info(item)
else:
# Do nothing (bad formatting)
item_type = item["item_type"]
_LOGGER.error(f"Invalid item_type: {item_type}")
pass
_LOGGER.info("Environment configuration loaded")
@@ -738,7 +733,7 @@ class Primaite(Env):
node_type,
node_priority,
node_hardware_state,
self.config_values,
self.training_config,
)
elif node_class == "ACTIVE":
# Active nodes have IP address, Software State and file system state
@@ -754,7 +749,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
elif node_class == "SERVICE":
# Service nodes have IP address, Software State, file system state and list of services
@@ -770,7 +765,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
node_services = item["services"]
for service in node_services:
@@ -849,14 +844,14 @@ class Primaite(Env):
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
ier_start_step = item["start_step"]
ier_end_step = item["end_step"]
ier_load = item["load"]
ier_protocol = item["protocol"]
ier_port = item["port"]
ier_source = item["source"]
ier_destination = item["destination"]
ier_mission_criticality = item["missionCriticality"]
ier_mission_criticality = item["mission_criticality"]
# Create IER and add to green IER dictionary
self.green_iers[ier_id] = IER(
@@ -879,14 +874,14 @@ class Primaite(Env):
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
ier_start_step = item["start_step"]
ier_end_step = item["end_step"]
ier_load = item["load"]
ier_protocol = item["protocol"]
ier_port = item["port"]
ier_source = item["source"]
ier_destination = item["destination"]
ier_mission_criticality = item["missionCriticality"]
ier_mission_criticality = item["mission_criticality"]
# Create IER and add to red IER dictionary
self.red_iers[ier_id] = IER(
@@ -909,8 +904,8 @@ class Primaite(Env):
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_start_step = item["start_step"]
pol_end_step = item["end_step"]
pol_node = item["nodeId"]
pol_type = NodePOLType[item["type"]]
@@ -943,8 +938,8 @@ class Primaite(Env):
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_start_step = item["start_step"]
pol_end_step = item["end_step"]
pol_target_node_id = item["targetNodeId"]
pol_initiator = NodePOLInitiator[item["initiator"]]
pol_type = NodePOLType[item["type"]]
@@ -1004,7 +999,7 @@ class Primaite(Env):
Args:
item: A config data item representing the services
"""
service_list = services["serviceList"]
service_list = services["service_list"]
for service in service_list:
service_name = service["name"]
@@ -1020,7 +1015,7 @@ class Primaite(Env):
Args:
item: A config data item representing the ports
"""
ports_list = ports["portsList"]
ports_list = ports["ports_list"]
for port in ports_list:
port_value = port["port"]
@@ -1029,6 +1024,14 @@ class Primaite(Env):
# Set the number of ports
self.num_ports = len(self.ports_list)
def get_observation_info(self, observation_info):
"""Extracts observation_info.
:param observation_info: Config item that defines which type of observation space to use
:type observation_info: str
"""
self.observation_type = ObservationType[observation_info["type"]]
def get_action_info(self, action_info):
"""
Extracts action_info.
@@ -1050,16 +1053,6 @@ class Primaite(Env):
"""
self.obs_config = obs_config
def get_steps_info(self, steps_info):
"""
Extracts steps_info.
Args:
item: A config data item representing steps info
"""
self.episode_steps = int(steps_info["steps"])
_LOGGER.info("Training episodes have " + str(self.episode_steps) + " steps")
def reset_environment(self):
"""
# Resets environment.
@@ -1067,11 +1060,11 @@ class Primaite(Env):
Uses config data config data in order to build the environment
configuration.
"""
for item in self.config_data:
if item["itemType"] == "NODE":
for item in self.lay_down_config:
if item["item_type"] == "NODE":
# Reset a node's state (normal and reference)
self.reset_node(item)
elif item["itemType"] == "ACL_RULE":
elif item["item_type"] == "ACL_RULE":
# Create an ACL rule (these are cleared on reset, so just need to recreate them)
self.create_acl_rule(item)
else:
@@ -1173,7 +1166,7 @@ class Primaite(Env):
def create_acl_action_dict(self):
"""Creates a dictionary mapping each possible discrete action to more readable multidiscrete action."""
# reserve 0 action to be a nothing action
actions = {0: [0, 0, 0, 0, 0, 0, 0]}
actions = {0: [0, 0, 0, 0, 0, 0]}
action_key = 1
# 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE
@@ -1185,16 +1178,14 @@ class Primaite(Env):
for dest_ip in range(self.num_nodes + 1):
for protocol in range(self.num_services + 1):
for port in range(self.num_ports + 1):
for position in range(self.max_acl_rules - 1):
action = [
action_decision,
action_permission,
source_ip,
dest_ip,
protocol,
port,
position,
]
action = [
action_decision,
action_permission,
source_ip,
dest_ip,
protocol,
port,
]
# Check to see if its an action we want to include as possible i.e. not a nothing action
if is_valid_acl_action_extra(action):
actions[action_key] = action

514
src/primaite/main.py
View File

@@ -1,29 +1,42 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Primaite - main (harness) module.
The main PrimAITE session runner module.
Coding Standards: PEP 8
TODO: This will eventually be refactored out into a proper Session class.
TODO: The passing about of session_dir and timestamp_str is temporary and
will be cleaned up once we move to a proper Session class.
"""
import logging
import os.path
import argparse
import json
import time
from datetime import datetime
from pathlib import Path
from typing import Final, Union
from uuid import uuid4
import yaml
from stable_baselines3 import A2C, PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
from stable_baselines3.ppo import MlpPolicy as PPOMlp
from primaite.common.config_values_main import ConfigValuesMain
from primaite import SESSIONS_DIR, getLogger
from primaite.config.training_config import TrainingConfig
from primaite.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import write_transaction_to_file
from primaite.transactions.transactions_to_file import \
write_transaction_to_file
# FUNCTIONS #
_LOGGER = getLogger(__name__)
def run_generic():
"""Run against a generic agent."""
def run_generic(env: Primaite, config_values: TrainingConfig):
"""
Run against a generic agent.
:param env: An instance of
:class:`~primaite.environment.primaite_env.Primaite`.
:param config_values: An instance of
:class:`~primaite.config.training_config.TrainingConfig`.
"""
for episode in range(0, config_values.num_episodes):
env.reset()
for step in range(0, config_values.num_steps):
@@ -47,9 +60,21 @@ def run_generic():
env.close()
def run_stable_baselines3_ppo():
"""Run against a stable_baselines3 PPO agent."""
if config_values.load_agent == True:
def run_stable_baselines3_ppo(
env: Primaite, config_values: TrainingConfig, session_path: Path, timestamp_str: str
):
"""
Run against a stable_baselines3 PPO agent.
:param env: An instance of
:class:`~primaite.environment.primaite_env.Primaite`.
:param config_values: An instance of
:class:`~primaite.config.training_config.TrainingConfig`.
: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>.
"""
if config_values.load_agent:
try:
agent = PPO.load(
config_values.agent_load_file,
@@ -62,30 +87,42 @@ def run_stable_baselines3_ppo():
"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)
_LOGGER.error("Could not load agent")
_LOGGER.error("Exception occured", exc_info=True)
else:
agent = PPO(PPOMlp, env, verbose=0, n_steps=config_values.num_steps)
if config_values.session_type == "TRAINING":
# We're in a training session
print("Starting training session...")
logging.info("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
save_agent(agent)
_LOGGER.debug("Starting training session...")
for episode in range(config_values.num_episodes):
agent.learn(total_timesteps=config_values.num_steps)
_save_agent(agent, session_path, timestamp_str)
else:
# Default to being in an evaluation session
print("Starting evaluation session...")
logging.info("Starting evaluation session...")
_LOGGER.debug("Starting evaluation session...")
evaluate_policy(agent, env, n_eval_episodes=config_values.num_episodes)
env.close()
def run_stable_baselines3_a2c():
"""Run against a stable_baselines3 A2C agent."""
if config_values.load_agent == True:
def run_stable_baselines3_a2c(
env: Primaite, config_values: TrainingConfig, session_path: Path, timestamp_str: str
):
"""
Run against a stable_baselines3 A2C agent.
:param env: An instance of
:class:`~primaite.environment.primaite_env.Primaite`.
:param config_values: An instance of
:class:`~primaite.config.training_config.TrainingConfig`.
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>.
"""
if config_values.load_agent:
try:
agent = A2C.load(
config_values.agent_load_file,
@@ -98,288 +135,219 @@ def run_stable_baselines3_a2c():
"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)
_LOGGER.error("Could not load agent")
_LOGGER.error("Exception occured", exc_info=True)
else:
agent = A2C("MlpPolicy", env, verbose=0, n_steps=config_values.num_steps)
if config_values.session_type == "TRAINING":
# We're in a training session
print("Starting training session...")
logging.info("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
save_agent(agent)
_LOGGER.debug("Starting training session...")
for episode in range(config_values.num_episodes):
agent.learn(total_timesteps=config_values.num_steps)
_save_agent(agent, session_path, timestamp_str)
else:
# Default to being in an evaluation session
print("Starting evaluation session...")
logging.info("Starting evaluation session...")
_LOGGER.debug("Starting evaluation session...")
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
time = now.strftime("%Y%m%d_%H%M%S")
def _write_session_metadata_file(
session_dir: Path, uuid: str, session_timestamp: datetime, env: Primaite
):
"""
Write the ``session_metadata.json`` file.
try:
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:
logging.error("Could not save agent")
logging.error("Exception occured", exc_info=True)
Creates a ``session_metadata.json`` in the ``session_dir`` directory
and adds the following key/value pairs:
- uuid: The UUID assigned to the session upon instantiation.
- start_datetime: The date & time the session started in iso format.
- end_datetime: NULL.
- total_episodes: NULL.
- total_time_steps: NULL.
- env:
- training_config:
- All training config items
- lay_down_config:
- All lay down config items
"""
metadata_dict = {
"uuid": uuid,
"start_datetime": session_timestamp.isoformat(),
"end_datetime": None,
"total_episodes": None,
"total_time_steps": None,
"env": {
"training_config": env.training_config.to_dict(json_serializable=True),
"lay_down_config": env.lay_down_config,
},
}
filepath = session_dir / "session_metadata.json"
_LOGGER.debug(f"Writing Session Metadata file: {filepath}")
with open(filepath, "w") as file:
json.dump(metadata_dict, file)
def configure_logging():
"""Configures logging."""
try:
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
filename = "logs/app_" + time + ".log"
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 Exception:
print("ERROR: Could not start logging")
def _update_session_metadata_file(session_dir: Path, env: Primaite):
"""
Update the ``session_metadata.json`` file.
Updates the `session_metadata.json`` in the ``session_dir`` directory
with the following key/value pairs:
- end_datetime: NULL.
- total_episodes: NULL.
- total_time_steps: NULL.
"""
with open(session_dir / "session_metadata.json", "r") as file:
metadata_dict = json.load(file)
metadata_dict["end_datetime"] = datetime.now().isoformat()
metadata_dict["total_episodes"] = env.episode_count
metadata_dict["total_time_steps"] = env.total_step_count
filepath = session_dir / "session_metadata.json"
_LOGGER.debug(f"Updating Session Metadata file: {filepath}")
with open(filepath, "w") as file:
json.dump(metadata_dict, file)
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"]
if "observationSpace" in config_data:
config_values.observation_config = config_data["observationSpace"]
else:
config_values.observation_config = None
config_values.num_episodes = int(config_data["numEpisodes"])
config_values.time_delay = int(config_data["timeDelay"])
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"]
# Environment
config_values.observation_space_high_value = int(
config_data["observationSpaceHighValue"]
)
# Reward values
# Generic
config_values.all_ok = int(config_data["allOk"])
# Node Hardware 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"])
# Node Software 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"])
# 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"])
# IER status
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"]
)
def _save_agent(agent: OnPolicyAlgorithm, session_path: Path, timestamp_str: str):
"""
Persist an agent.
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"
)
Only works for stable baselines3 agents at present.
except Exception:
logging.error("Could not save load config data")
logging.error("Exception occured", exc_info=True)
: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>.
"""
if not isinstance(agent, OnPolicyAlgorithm):
msg = f"Can only save {OnPolicyAlgorithm} agents, got {type(agent)}."
_LOGGER.error(msg)
else:
filepath = session_path / f"agent_saved_{timestamp_str}"
agent.save(filepath)
_LOGGER.debug(f"Trained agent saved as: {filepath}")
# MAIN PROCESS #
def _get_session_path(session_timestamp: datetime) -> Path:
"""
Get the directory path the session will output to.
# Starting point
This is set in the format of:
~/primaite/sessions/<yyyy-mm-dd>/<yyyy-mm-dd>_<hh-mm-ss>.
# Welcome message
print("Welcome to the Primary-level AI Training Environment (PrimAITE)")
:param session_timestamp: This is the datetime that the session started.
:return: The session directory path.
"""
date_dir = session_timestamp.strftime("%Y-%m-%d")
session_dir = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
session_path = SESSIONS_DIR / date_dir / session_dir
session_path.mkdir(exist_ok=True, parents=True)
# Configure logging
configure_logging()
return session_path
# Open the main config file
try:
config_file_main = open("config/config_main.yaml", "r")
config_data = yaml.safe_load(config_file_main)
# Create a config class
config_values = ConfigValuesMain()
# Load in config data
load_config_values()
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 #
# - initial observation space
# - action
# - reward
# - new observation space
transaction_list = []
def run(training_config_path: Union[str, Path], lay_down_config_path: Union[str, Path]):
"""Run the PrimAITE Session.
# Create the Primaite environment
# try:
env = Primaite(config_values, transaction_list)
# logging.info("PrimAITE environment created")
# except Exception:
# logging.error("Could not create PrimAITE environment")
# logging.error("Exception occured", exc_info=True)
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.
"""
# Welcome message
print("Welcome to the Primary-level AI Training Environment (PrimAITE)")
uuid = str(uuid4())
session_timestamp: Final[datetime] = datetime.now()
session_dir = _get_session_path(session_timestamp)
timestamp_str = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = env.episode_steps
print(f"The output directory for this session is: {session_dir}")
# Run environment against an agent
if config_values.agent_identifier == "GENERIC":
run_generic()
elif config_values.agent_identifier == "STABLE_BASELINES3_PPO":
run_stable_baselines3_ppo()
elif config_values.agent_identifier == "STABLE_BASELINES3_A2C":
run_stable_baselines3_a2c()
# Create a list of transactions
# A transaction is an object holding the:
# - episode #
# - step #
# - initial observation space
# - action
# - reward
# - new observation space
transaction_list = []
print("Session finished")
logging.info("Session finished")
# Create the Primaite environment
env = Primaite(
training_config_path=training_config_path,
lay_down_config_path=lay_down_config_path,
transaction_list=transaction_list,
session_path=session_dir,
timestamp_str=timestamp_str,
)
print("Saving transaction logs...")
logging.info("Saving transaction logs...")
print("Writing Session Metadata file...")
write_transaction_to_file(transaction_list)
_write_session_metadata_file(
session_dir=session_dir, uuid=uuid, session_timestamp=session_timestamp, env=env
)
config_values = env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = env.episode_steps
# Run environment against an agent
if config_values.agent_identifier == "GENERIC":
run_generic(env=env, config_values=config_values)
elif config_values.agent_identifier == "STABLE_BASELINES3_PPO":
run_stable_baselines3_ppo(
env=env,
config_values=config_values,
session_path=session_dir,
timestamp_str=timestamp_str,
)
elif config_values.agent_identifier == "STABLE_BASELINES3_A2C":
run_stable_baselines3_a2c(
env=env,
config_values=config_values,
session_path=session_dir,
timestamp_str=timestamp_str,
)
print("Session finished")
_LOGGER.debug("Session finished")
print("Saving transaction logs...")
write_transaction_to_file(
transaction_list=transaction_list,
session_path=session_dir,
timestamp_str=timestamp_str,
)
print("Updating Session Metadata file...")
_update_session_metadata_file(session_dir=session_dir, env=env)
print("Finished")
_LOGGER.debug("Finished")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tc")
parser.add_argument("--ldc")
args = parser.parse_args()
if not args.tc:
_LOGGER.error(
"Please provide a training config file using the --tc " "argument"
)
if not args.ldc:
_LOGGER.error(
"Please provide a lay down config file using the --ldc " "argument"
)
run(training_config_path=args.tc, lay_down_config_path=args.ldc)
config_file_main.close()
print("Finished")
logging.info("Finished")

18
src/primaite/nodes/active_node.py
View File

@@ -3,7 +3,6 @@
import logging
from typing import Final
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import (
FileSystemState,
HardwareState,
@@ -11,6 +10,7 @@ from primaite.common.enums import (
Priority,
SoftwareState,
)
from primaite.config.training_config import TrainingConfig
from primaite.nodes.node import Node
_LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
@@ -29,7 +29,7 @@ class ActiveNode(Node):
ip_address: str,
software_state: SoftwareState,
file_system_state: FileSystemState,
config_values: ConfigValuesMain,
config_values: TrainingConfig,
):
"""
Init.
@@ -210,3 +210,17 @@ class ActiveNode(Node):
self.file_system_state_observed = self.file_system_state_actual
self.file_system_scanning = False
self.file_system_scanning_count = 0
def update_resetting_status(self):
"""Updates the reset count & makes software and file state to GOOD."""
super().update_resetting_status()
if self.resetting_count <= 0:
self.file_system_state_actual = FileSystemState.GOOD
self.software_state = SoftwareState.GOOD
def update_booting_status(self):
"""Updates the booting software and file state to GOOD."""
super().update_booting_status()
if self.booting_count <= 0:
self.file_system_state_actual = FileSystemState.GOOD
self.software_state = SoftwareState.GOOD

27
src/primaite/nodes/node.py
View File

@@ -2,8 +2,8 @@
"""The base Node class."""
from typing import Final
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import HardwareState, NodeType, Priority
from primaite.config.training_config import TrainingConfig
class Node:
@@ -16,7 +16,7 @@ class Node:
node_type: NodeType,
priority: Priority,
hardware_state: HardwareState,
config_values: ConfigValuesMain,
config_values: TrainingConfig,
):
"""
Init.
@@ -34,7 +34,9 @@ class Node:
self.priority = priority
self.hardware_state: HardwareState = hardware_state
self.resetting_count: int = 0
self.config_values: ConfigValuesMain = config_values
self.config_values: TrainingConfig = config_values
self.booting_count: int = 0
self.shutting_down_count: int = 0
def __repr__(self):
"""Returns the name of the node."""
@@ -42,11 +44,12 @@ class Node:
def turn_on(self):
"""Sets the node state to ON."""
self.hardware_state = HardwareState.ON
self.hardware_state = HardwareState.BOOTING
self.booting_count = self.config_values.node_booting_duration
def turn_off(self):
"""Sets the node state to OFF."""
self.hardware_state = HardwareState.OFF
self.shutting_down_count = self.config_values.node_shutdown_duration
def reset(self):
"""Sets the node state to Resetting and starts the reset count."""
@@ -59,3 +62,17 @@ class Node:
if self.resetting_count <= 0:
self.resetting_count = 0
self.hardware_state = HardwareState.ON
def update_booting_status(self):
"""Updates the booting count"""
self.booting_count -= 1
if self.booting_count <= 0:
self.booting_count = 0
self.hardware_state = HardwareState.ON
def update_shutdown_status(self):
"""Updates the shutdown count"""
self.shutting_down_count -= 1
if self.shutting_down_count <= 0:
self.shutting_down_count = 0
self.hardware_state = HardwareState.OFF

4
src/primaite/nodes/passive_node.py
View File

@@ -1,7 +1,7 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The Passive Node class (i.e. an actuator)."""
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import HardwareState, NodeType, Priority
from primaite.config.training_config import TrainingConfig
from primaite.nodes.node import Node
@@ -15,7 +15,7 @@ class PassiveNode(Node):
node_type: NodeType,
priority: Priority,
hardware_state: HardwareState,
config_values: ConfigValuesMain,
config_values: TrainingConfig,
):
"""
Init.

16
src/primaite/nodes/service_node.py
View File

@@ -3,7 +3,6 @@
import logging
from typing import Dict, Final
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import (
FileSystemState,
HardwareState,
@@ -12,6 +11,7 @@ from primaite.common.enums import (
SoftwareState,
)
from primaite.common.service import Service
from primaite.config.training_config import TrainingConfig
from primaite.nodes.active_node import ActiveNode
_LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
@@ -30,7 +30,7 @@ class ServiceNode(ActiveNode):
ip_address: str,
software_state: SoftwareState,
file_system_state: FileSystemState,
config_values: ConfigValuesMain,
config_values: TrainingConfig,
):
"""
Init.
@@ -188,3 +188,15 @@ class ServiceNode(ActiveNode):
for service_key, service_value in self.services.items():
if service_value.software_state == SoftwareState.PATCHING:
service_value.reduce_patching_count()
def update_resetting_status(self):
super().update_resetting_status()
if self.resetting_count <= 0:
for service in self.services.values():
service.software_state = SoftwareState.GOOD
def update_booting_status(self):
super().update_booting_status()
if self.booting_count <= 0:
for service in self.services.values():
service.software_state =SoftwareState.GOOD

32
src/primaite/notebooks/__init__.py Normal file
View File

@@ -0,0 +1,32 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import importlib.util
import os
import subprocess
import sys
from primaite import NOTEBOOKS_DIR, getLogger
_LOGGER = getLogger(__name__)
def start_jupyter_session():
"""
Starts a new Jupyter notebook session in the app notebooks directory.
Currently only works on Windows OS.
.. todo:: Figure out how to get this working for Linux and MacOS too.
"""
if importlib.util.find_spec("jupyter") is not None:
jupyter_cmd = "python3 -m jupyter lab"
if sys.platform == "win32":
jupyter_cmd = "jupyter lab"
working_dir = os.getcwd()
os.chdir(NOTEBOOKS_DIR)
subprocess.Popen(jupyter_cmd)
os.chdir(working_dir)
else:
# Jupyter is not installed
_LOGGER.error("Cannot start jupyter lab as it is not installed")

1
src/primaite/setup/__init__.py Normal file
View File

@@ -0,0 +1 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

5
src/primaite/setup/_package_data/primaite_config.yaml Normal file
View File

@@ -0,0 +1,5 @@
# The main PrimAITE application config file
# Logging
log_level: INFO
logger_format: '%(asctime)s::%(levelname)s::%(name)s::%(lineno)s::%(message)s'

13
src/primaite/setup/old_installation_clean_up.py Normal file
View File

@@ -0,0 +1,13 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from primaite import getLogger
_LOGGER = getLogger(__name__)
def run():
"""Perform the full clean-up."""
pass
if __name__ == "__main__":
run()

39
src/primaite/setup/reset_demo_notebooks.py Normal file
View File

@@ -0,0 +1,39 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import filecmp
import os
import shutil
from pathlib import Path
import pkg_resources
from primaite import NOTEBOOKS_DIR, getLogger
_LOGGER = getLogger(__name__)
def run(overwrite_existing: bool = True):
"""
Resets the demo jupyter notebooks in the users app notebooks directory.
:param overwrite_existing: A bool to toggle replacing existing edited
notebooks on or off.
"""
notebooks_package_data_root = pkg_resources.resource_filename(
"primaite", "notebooks/_package_data"
)
for subdir, dirs, files in os.walk(notebooks_package_data_root):
for file in files:
fp = os.path.join(subdir, file)
path_split = os.path.relpath(fp, notebooks_package_data_root).split(os.sep)
target_fp = NOTEBOOKS_DIR / Path(*path_split)
target_fp.parent.mkdir(exist_ok=True, parents=True)
copy_file = not target_fp.is_file()
if overwrite_existing and not copy_file:
copy_file = (not filecmp.cmp(fp, target_fp)) and (
".ipynb_checkpoints" not in str(target_fp)
)
if copy_file:
shutil.copy2(fp, target_fp)
_LOGGER.info(f"Reset example notebook: {target_fp}")

37
src/primaite/setup/reset_example_configs.py Normal file
View File

@@ -0,0 +1,37 @@
import filecmp
import os
import shutil
from pathlib import Path
import pkg_resources
from primaite import USERS_CONFIG_DIR, getLogger
_LOGGER = getLogger(__name__)
def run(overwrite_existing=True):
"""
Resets the example config files in the users app config directory.
:param overwrite_existing: A bool to toggle replacing existing edited
config on or off.
"""
configs_package_data_root = pkg_resources.resource_filename(
"primaite", "config/_package_data"
)
for subdir, dirs, files in os.walk(configs_package_data_root):
for file in files:
fp = os.path.join(subdir, file)
path_split = os.path.relpath(fp, configs_package_data_root).split(os.sep)
target_fp = USERS_CONFIG_DIR / "example_config" / Path(*path_split)
target_fp.parent.mkdir(exist_ok=True, parents=True)
copy_file = not target_fp.is_file()
if overwrite_existing and not copy_file:
copy_file = not filecmp.cmp(fp, target_fp)
if copy_file:
shutil.copy2(fp, target_fp)
_LOGGER.info(f"Reset example config: {target_fp}")

27
src/primaite/setup/setup_app_dirs.py Normal file
View File

@@ -0,0 +1,27 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from primaite import _USER_DIRS, LOG_DIR, NOTEBOOKS_DIR, getLogger
_LOGGER = getLogger(__name__)
def run():
"""
Handles creation of application directories and user directories.
Uses `platformdirs.PlatformDirs` and `pathlib.Path` to create the required
app directories in the correct locations based on the users OS.
"""
app_dirs = [
_USER_DIRS,
NOTEBOOKS_DIR,
LOG_DIR,
]
for app_dir in app_dirs:
if not app_dir.is_dir():
app_dir.mkdir(parents=True, exist_ok=True)
_LOGGER.info(f"Created directory: {app_dir}")
if __name__ == "__main__":
run()

50
src/primaite/transactions/transactions_to_file.py
View File

@@ -2,9 +2,11 @@
"""Writes the Transaction log list out to file for evaluation to utilse."""
import csv
import logging
import os.path
from datetime import datetime
from pathlib import Path
from primaite import getLogger
_LOGGER = getLogger(__name__)
def turn_action_space_to_array(_action_space):
@@ -14,11 +16,10 @@ def turn_action_space_to_array(_action_space):
Args:
_action_space: The action space.
"""
return_array = []
for x in range(len(_action_space)):
return_array.append(str(_action_space[x]))
return return_array
if isinstance(_action_space, list):
return [str(i) for i in _action_space]
else:
return [str(_action_space)]
def turn_obs_space_to_array(_obs_space, _obs_assets, _obs_features):
@@ -38,18 +39,22 @@ def turn_obs_space_to_array(_obs_space, _obs_assets, _obs_features):
return return_array
def write_transaction_to_file(_transaction_list):
def write_transaction_to_file(transaction_list, session_path: Path, timestamp_str: str):
"""
Writes transaction logs to file to support training evaluation.
Args:
_transaction_list: The list of transactions from all steps and all episodes
_num_episodes: The number of episodes that were conducted.
:param transaction_list: The list of transactions from all steps and all
episodes.
: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>.
"""
# Get the first transaction and use it to determine the makeup of the observation space and action space
# Label the obs space fields in csv as "OSI_1_1", "OSN_1_1" and action space as "AS_1"
# Get the first transaction and use it to determine the makeup of the
# observation space and action space
# Label the obs space fields in csv as "OSI_1_1", "OSN_1_1" and action
# space as "AS_1"
# This will be tied into the PrimAITE Use Case so that they make sense
template_transation = _transaction_list[0]
template_transation = transaction_list[0]
action_length = template_transation.action_space.size
obs_shape = template_transation.obs_space_post.shape
obs_assets = template_transation.obs_space_post.shape[0]
@@ -75,21 +80,15 @@ def write_transaction_to_file(_transaction_list):
# Open up a csv file
header = ["Timestamp", "Episode", "Step", "Reward"]
header = header + action_header + obs_header_initial + obs_header_new
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
try:
path = "outputs/results/"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/results/all_transactions_" + time + ".csv"
filename = session_path / f"all_transactions_{timestamp_str}.csv"
_LOGGER.debug(f"Saving transaction logs: {filename}")
csv_file = open(filename, "w", encoding="UTF8", newline="")
csv_writer = csv.writer(csv_file)
csv_writer.writerow(header)
for transaction in _transaction_list:
for transaction in transaction_list:
csv_data = [
str(transaction.timestamp),
str(transaction.episode_number),
@@ -110,5 +109,4 @@ def write_transaction_to_file(_transaction_list):
csv_file.close()
except Exception:
logging.error("Could not save the transaction file")
logging.error("Exception occured", exc_info=True)
_LOGGER.error("Could not save the transaction file", exc_info=True)

0
src/primaite/utils/__init__.py Normal file
View File

32
src/primaite/utils/package_data.py Normal file
View File

@@ -0,0 +1,32 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import os
from pathlib import Path
import pkg_resources
from primaite import getLogger
_LOGGER = getLogger(__name__)
def get_file_path(path: str) -> Path:
"""
Get PrimAITE package data.
:Example:
>>> from primaite.utils.package_data import get_file_path
>>> main_env_config = get_file_path("config/_package_data/training_config_main.yaml")
:param path: The path from the primaite root.
:return: The file path of the package data file.
:raise FileNotFoundError: When the filepath does not exist.
"""
fp = pkg_resources.resource_filename("primaite", path)
if os.path.isfile(fp):
return Path(fp)
else:
msg = f"Cannot PrimAITE package data: {fp}"
_LOGGER.error(msg)
raise FileNotFoundError(msg)