Initial commit of v1.0.0. Updated the .gitignore for the standard Python gitignore. Added Azure DevOps release pipeline for proper artifact release from the start.

This commit is contained in:
Chris McCarthy
2023-03-28 17:33:34 +01:00
parent fdbd5903bc
commit 7800f1f66e
45 changed files with 5086 additions and 20 deletions

6
.azure/.pypirc Normal file
View File

@@ -0,0 +1,6 @@
[distutils]
Index-servers =
PrimAITE
[PrimAITE]
Repository = https://pkgs.dev.azure.com/ma-dev-uk/PrimAITE/_packaging/PrimAITE/pypi/upload/

View File

@@ -0,0 +1,37 @@
trigger:
- main
pool:
vmImage: ubuntu-latest
strategy:
matrix:
python.version: '3.10'
steps:
- task: UsePythonVersion@0
inputs:
versionSpec: '$(python.version)'
displayName: 'Use Python $(python.version)'
- script: |
python -m pip install --upgrade pip
pip install build
pip install wheel
pip install twine
pip install keyring
pip install artifacts-keyring
displayName: 'Install build dependencies'
- script: |
python setup.py sdist bdist_wheel
displayName: 'Build PrimAITE sdist and wheel'
- task: TwineAuthenticate@1
displayName: 'Twine Authenticate'
inputs:
artifactFeed: PrimAITE/PrimAITE
- script: |
python -m twine upload --verbose -r PrimAITE --config-file $(PYPIRC_PATH) dist/*
displayName: 'Artifact Upload'

143
.gitignore vendored Normal file
View File

@@ -0,0 +1,143 @@
# PrimAITE Package
PRIMAITE/outputs
PRIMAITE/outputs/*
TestResults
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/

261
PRIMAITE/Main.py Normal file
View File

@@ -0,0 +1,261 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
PRIMAITE - main (harness) module
Coding Standards: PEP 8
"""
from sys import exc_info
import time
import yaml
import os.path
import logging
from datetime import datetime
from environment.primaite import PRIMAITE
from transactions.transactions_to_file import write_transaction_to_file
from common.config_values_main import config_values_main
from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy as PPOMlp
from stable_baselines3 import A2C
from stable_baselines3.common.env_checker import check_env
from stable_baselines3.common.evaluation import evaluate_policy
################################# FUNCTIONS ######################################
def run_generic():
"""
Run against a generic agent
"""
for episode in range(0, config_values.num_episodes):
for step in range(0, config_values.num_steps):
# Send the observation space to the agent to get an action
# TEMP - random action for now
# action = env.blue_agent_action(obs)
action = env.action_space.sample()
# Run the simulation step on the live environment
obs, reward, done, info = env.step(action)
# Break if done is True
if done:
break
# Introduce a delay between steps
time.sleep(config_values.time_delay / 1000)
# Reset the environment at the end of the episode
env.reset()
env.close()
def run_stable_baselines3_ppo():
"""
Run against a stable_baselines3 PPO agent
"""
#if check_env(env, warn=TRUE):
# print("Environment is NOT OpenAI Gym Compliant")
#else:
# print("Environment is OpenAI Gym Compliant")
agent = PPO(PPOMlp, env, verbose=0, n_steps=config_values.num_steps)
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
env.close()
save_agent(agent)
def run_stable_baselines3_a2c():
"""
Run against a stable_baselines3 A2C agent
"""
#if check_env(env, warn=TRUE):
# print("Environment is NOT OpenAI Gym Compliant")
#else:
# print("Environment is OpenAI Gym Compliant")
agent = A2C("MlpPolicy", env, verbose=0, n_steps=config_values.num_steps)
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
env.close()
save_agent(agent)
def save_agent(_agent):
"""
Persist an agent (only works for stable baselines3 agents at present)
"""
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
try:
path = 'outputs/agents/'
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/agents/agent_saved_" + time
_agent.save(filename)
logging.info("Trained agent saved as " + filename)
except Exception as e:
logging.error("Could not save agent")
logging.error("Exception occured", exc_info=True)
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:
print("ERROR: Could not start logging")
def load_config_values():
"""
Loads the config values from the main config file into a config object
"""
try:
# Generic
config_values.agent_identifier = config_data['agentIdentifier']
config_values.num_episodes = int(config_data['numEpisodes'])
config_values.time_delay = int(config_data['timeDelay'])
config_values.config_filename_use_case = config_data['configFilename']
# Environment
config_values.observation_space_high_value = int(config_data['observationSpaceHighValue'])
# Reward values
# Generic
config_values.all_ok = int(config_data['allOk'])
# Node Operating State
config_values.off_should_be_on = int(config_data['offShouldBeOn'])
config_values.off_should_be_resetting = int(config_data['offShouldBeResetting'])
config_values.on_should_be_off = int(config_data['onShouldBeOff'])
config_values.on_should_be_resetting = int(config_data['onShouldBeResetting'])
config_values.resetting_should_be_on = int(config_data['resettingShouldBeOn'])
config_values.resetting_should_be_off = int(config_data['resettingShouldBeOff'])
# Node O/S or Service State
config_values.good_should_be_patching = int(config_data['goodShouldBePatching'])
config_values.good_should_be_compromised = int(config_data['goodShouldBeCompromised'])
config_values.good_should_be_overwhelmed = int(config_data['goodShouldBeOverwhelmed'])
config_values.patching_should_be_good = int(config_data['patchingShouldBeGood'])
config_values.patching_should_be_compromised = int(config_data['patchingShouldBeCompromised'])
config_values.patching_should_be_overwhelmed = int(config_data['patchingShouldBeOverwhelmed'])
config_values.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'])
# 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'])
logging.info("Training agent: " + config_values.agent_identifier)
logging.info("Training environment config: " + config_values.config_filename_use_case)
logging.info("Training cycle has " + str(config_values.num_episodes) + " episodes")
except Exception as e:
logging.error("Could not save load config data")
logging.error("Exception occured", exc_info=True)
################################# MAIN PROCESS ############################################
# Starting point
# Welcome message
print("Welcome to the Primary-level AI Training Environment (PrimAITE)")
# Configure logging
configure_logging()
# 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 = config_values_main()
# Load in config data
load_config_values()
except Exception as e:
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 = []
# Create the PRIMAITE environment
try:
env = PRIMAITE(config_values, transaction_list)
logging.info("PrimAITE environment created")
except Exception as e:
logging.error("Could not create PrimAITE environment")
logging.error("Exception occured", exc_info=True)
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = env.episode_steps
print("Starting training...")
logging.info("Training started...")
# 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()
print("Finished training")
logging.info("Training complete")
print("Saving transaction logs...")
logging.info("Saving transaction logs...")
write_transaction_to_file(transaction_list)
config_file_main.close
print("Finished")
logging.info("Finished")

1
PRIMAITE/acl/__init__.py Normal file
View File

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

View File

@@ -0,0 +1,134 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
A class that implements the access control list implementation for the network
"""
from acl.acl_rule import ACLRule
class AccessControlList():
"""
Access Control List class
"""
def __init__(self):
"""
Init
"""
self.acl = {} # A dictionary of ACL Rules
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
"""
Checks for IP address matches
Args:
_rule: The rule being checked
_source_ip_address: the source IP address to compare
_dest_ip_address: the destination IP address to compare
Returns:
True if match; False otherwise.
"""
if ((_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == _dest_ip_address) or
(_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == _dest_ip_address) or
(_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == "ANY") or
(_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == "ANY")):
return True
else:
return False
def is_blocked(self, _source_ip_address, _dest_ip_address, _protocol, _port):
"""
Checks for rules that block a protocol / port
Args:
_source_ip_address: the source IP address to check
_dest_ip_address: the destination IP address to check
_protocol: the protocol to check
_port: the port to check
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):
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")):
# There's a matching rule. Get the permission
if rule_value.get_permission() == "DENY":
return True
elif rule_value.get_permission() == "ALLOW":
return False
# If there has been no rule to allow the IER through, it will return a blocked signal by default
return True
def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Adds a new rule
Args:
_permission: the permission value (e.g. "ALLOW" or "DENY")
_source_ip: the source IP address
_dest_ip: the destination IP address
_protocol: the protocol
_port: the port
"""
new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
hash_value = hash(new_rule)
self.acl[hash_value] = new_rule
def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Removes a rule
Args:
_permission: the permission value (e.g. "ALLOW" or "DENY")
_source_ip: the source IP address
_dest_ip: the destination IP address
_protocol: the protocol
_port: the port
"""
rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
hash_value = hash(rule)
# There will not always be something 'popable' since the agent will be trying random things
try:
self.acl.pop(hash_value)
except:
return
def remove_all_rules(self):
"""
Removes all rules
"""
self.acl.clear()
def get_dictionary_hash(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Produces a hash value for a rule
Args:
_permission: the permission value (e.g. "ALLOW" or "DENY")
_source_ip: the source IP address
_dest_ip: the destination IP address
_protocol: the protocol
_port: the port
Returns:
Hash value based on rule parameters.
"""
rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
hash_value = hash(rule)
return hash_value

88
PRIMAITE/acl/acl_rule.py Normal file
View File

@@ -0,0 +1,88 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
A class that implements an access control list rule
"""
class ACLRule():
"""
Access Control List Rule class
"""
def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Init
Args:
_permission: The permission (ALLOW or DENY)
_source_ip: The source IP address
_dest_ip: The destination IP address
_protocol: The rule protocol
_port: The rule port
"""
self.permission = _permission
self.source_ip = _source_ip
self.dest_ip = _dest_ip
self.protocol = _protocol
self.port = _port
def __hash__(self):
"""
Override the hash function
Returns:
Returns hash of core parameters.
"""
return hash((self.permission, self.source_ip, self.dest_ip, self.protocol, self.port))
def get_permission(self):
"""
Gets the permission attribute
Returns:
Returns permission attribute
"""
return self.permission
def get_source_ip(self):
"""
Gets the source IP address attribute
Returns:
Returns source IP address attribute
"""
return self.source_ip
def get_dest_ip(self):
"""
Gets the desintation IP address attribute
Returns:
Returns destination IP address attribute
"""
return self.dest_ip
def get_protocol(self):
"""
Gets the protocol attribute
Returns:
Returns protocol attribute
"""
return self.protocol
def get_port(self):
"""
Gets the port attribute
Returns:
Returns port attribute
"""
return self.port

View File

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

View File

@@ -0,0 +1,59 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The config class
"""
class config_values_main(object):
"""
Class to hold main config values
"""
def __init__(self):
"""
Init
"""
# Generic
self.agent_identifier = "" # the agent in use
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
# Environment
self.observation_space_high_value = 0 # The high value for the observation space
# Reward values
# Generic
self.all_ok = 0
# Node Operating 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
# Node O/S 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.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
# 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

84
PRIMAITE/common/enums.py Normal file
View File

@@ -0,0 +1,84 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Enumerations for APE
"""
from enum import Enum
class TYPE(Enum):
"""
Node type enumeration
"""
CCTV = 1
SWITCH = 2
COMPUTER = 3
LINK = 4
MONITOR = 5
PRINTER = 6
LOP = 7
RTU = 8
ACTUATOR = 9
SERVER = 10
class PRIORITY(Enum):
"""
Node priority enumeration
"""
P1 = 1
P2 = 2
P3 = 3
P4 = 4
P5 = 5
class HARDWARE_STATE(Enum):
"""
Node hardware state enumeration
"""
ON = 1
OFF = 2
RESETTING = 3
class SOFTWARE_STATE(Enum):
"""
O/S or Service state enumeration
"""
GOOD = 1
PATCHING = 2
COMPROMISED = 3
OVERWHELMED = 4
class NODE_POL_TYPE(Enum):
"""
Node Pattern of Life type enumeration
"""
OPERATING = 1
OS = 2
SERVICE = 3
class PROTOCOL(Enum):
"""
Service protocol enumeration
"""
LDAP = 0
FTP = 1
HTTPS = 2
SMTP = 3
RTP = 4
IPP = 5
TCP = 6
NONE = 7
class ACTION_TYPE(Enum):
"""
Action type enumeration
"""
NODE = 0
ACL = 1

View File

@@ -0,0 +1,59 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The protocol class
"""
class Protocol(object):
"""
Protocol class
"""
def __init__(self, _name):
"""
Init
Args:
_name: The protocol name
"""
self.name = _name
self.load = 0 # bps
def get_name(self):
"""
Gets the protocol name
Returns:
The protocol name
"""
return self.name
def get_load(self):
"""
Gets the protocol load
Returns:
The protocol load (bps)
"""
return self.load
def add_load(self, _load):
"""
Adds load to the protocol
Args:
_load: The load to add
"""
self.load += _load
def clear_load(self):
"""
Clears the load on this protocol
"""
self.load = 0

100
PRIMAITE/common/service.py Normal file
View File

@@ -0,0 +1,100 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The Service class
"""
from common.enums import SOFTWARE_STATE
class Service(object):
"""
Service class
"""
def __init__(self, _name, _port, _state):
"""
Init
Args:
_name: The service name
_port: The service port
_state: The service state
"""
self.name = _name
self.port = _port
self.state = _state
self.patching_count = 0
def set_name(self, _name):
"""
Sets the service name
Args:
_name: The service name
"""
self.name = _name
def get_name(self):
"""
Gets the service name
Returns:
The service name
"""
return self.name
def set_port(self, _port):
"""
Sets the service port
Args:
_port: The service port
"""
self.port = _port
def get_port(self):
"""
Gets the service port
Returns:
The service port
"""
return self.port
def set_state(self, _state):
"""
Sets the service state
Args:
_state: The service state
"""
self.state = _state
def get_state(self):
"""
Gets the service state
Returns:
The service state
"""
return self.state
def reduce_patching_count(self):
"""
Reduces the patching count for the service
"""
self.patching_count -= 1
if self.patching_count <= 0:
self.patching_count = 0
self.state = SOFTWARE_STATE.GOOD

View File

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

View File

@@ -0,0 +1,52 @@
# 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_PPO
# Maximum number of episodes to run per training session
numEpisodes: 10
# Time delay between steps (for generic agents)
timeDelay: 10
# Filename of the scenario / laydown
configFilename: config_2_DDOS_BASIC.yaml
# Environment config values
# The high value for the observation space
observationSpaceHighValue: 1000000000
# Reward values
# Generic
allOk: 0
# Node Operating State
offShouldBeOn: -10
offShouldBeResetting: -5
onShouldBeOff: -2
onShouldBeResetting: -5
resettingShouldBeOn: -5
resettingShouldBeOff: -2
# Node O/S or Service State
goodShouldBePatching: 2
goodShouldBeCompromised: 5
goodShouldBeOverwhelmed: 5
patchingShouldBeGood: -5
patchingShouldBeCompromised: 2
patchingShouldBeOverwhelmed: 2
compromisedShouldBeGood: -20
compromisedShouldBePatching: -20
compromisedShouldBeOverwhelmed: -20
compromised: -20
overwhelmedShouldBeGood: -20
overwhelmedShouldBePatching: -20
overwhelmedShouldBeCompromised: -20
overwhelmed: -20
# 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

20
PRIMAITE/docs/Makefile Normal file
View File

@@ -0,0 +1,20 @@
# Minimal makefile for Sphinx documentation
#
# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS ?=
SPHINXBUILD ?= sphinx-build
SOURCEDIR = source
BUILDDIR = build
# Put it first so that "make" without argument is like "make help".
help:
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
.PHONY: help Makefile
# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

35
PRIMAITE/docs/make.bat Normal file
View File

@@ -0,0 +1,35 @@
@ECHO OFF
pushd %~dp0
REM Command file for Sphinx documentation
if "%SPHINXBUILD%" == "" (
set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=source
set BUILDDIR=build
%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
echo.
echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
echo.installed, then set the SPHINXBUILD environment variable to point
echo.to the full path of the 'sphinx-build' executable. Alternatively you
echo.may add the Sphinx directory to PATH.
echo.
echo.If you don't have Sphinx installed, grab it from
echo.https://www.sphinx-doc.org/
exit /b 1
)
if "%1" == "" goto help
%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end
:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
:end
popd

View File

@@ -0,0 +1,308 @@
.. _about:
About PrimAITE
==============
Features
********
PrimAITE provides the following features:
* A flexible network / system laydown based on the Python networkx framework
* Nodes and links (edges) host Python classes in order to present attributes and methods (and hence, a more representative model of a platform / system)
* A green agent Information Exchange Requirement (IER) function allows the representation of traffic (protocols and loading) on any / all links. Application of IERs is based on the status of node operating systems and services
* A green agent node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. Operating state, Operating System state, Service state)
* An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP, destination IP, protocol and port). Application of IERs adheres to any ACL restrictions
* Presents an OpenAI Gym interface to the environment, allowing integration with any OpenAI Gym compliant defensive agents
* Red agent activity based on red IERs and red PoL
* Defined reward function for use with RL agents (based on nodes status, and green / red IER success)
* Fully configurable (network / system laydown, IERs, node PoL, ACL, episode step period, episode max steps) and repeatable to suit the training requirements of agents. Therefore, not bound to a representation of any particular platform, system or technology
* Full capture of discrete metrics relating to agent training (full system state, agent actions taken, average reward)
* Networkx provides laydown visualisation capability
Architecture - Nodes and Links
******************************
**Nodes**
An inheritance model has been adopted in order to model nodes. All nodes have the following base attributes (Class: Node):
* ID
* Name
* Type (e.g. computer, switch, RTU - enumeration)
* Priority (P1, P2, P3, P4 or P5 - enumeration)
* Operating State (ON, OFF, RESETTING - enumeration)
Active Nodes also have the following attributes (Class: Active Node):
* IP Address
* Operating System State (GOOD, PATCHING, COMPROMISED - enumeration)
Service Nodes also have the following attributes (Class: Service Node):
* List of Services (where service is composed of service name and port). There is no theoretical limit on the number of services that can be modelled. Services and protocols are currently intrinsically linked (i.e. a service is an application on a node transmitting traffic of this protocol type)
* Service state (GOOD, PATCHING, COMPROMISED, OVERWHELMED - enumeration)
Passive Nodes are currently not used (but may be employed for non IP-based components such as machinery actuators in future releases).
**Links**
Links are modelled both as network edges (networkx) and as Python classes, in order to extend their functionality. Links include the following attributes:
* ID
* Name
* Bandwidth (bits/s)
* Source node ID
* Destination node ID
* Protocol list (containing the loading of protocols currently running on the link)
When the simulation runs, IERs are applied to the links in order to model traffic loading, individually assigned to each protocol. This allows green (background) and red agent behaviour to be modelled, and defensive agents to identify suspicious traffic patterns at a protocol / traffic loading level of fidelity.
Information Exchange Requirements (IERs)
****************************************
PrimAITE adopts the concept of Information Exchange Requirements (IERs) to model both green agent (background) and red agent (adversary) behaviour. IERs are used to initiate modelling of traffic loading on the network, and have the following attributes:
* ID
* Start step (i.e. which step in the training episode should the IER start)
* End step (i.e. which step in the training episode should the IER end)
* Source node ID
* Destination node ID
* Load (bits/s)
* Protocol
* Port
* Running status (i.e. on / off)
The application of green agent IERs between a source and destination follows a number of rules. Specifically:
1. Does the current simulation time step fall between IER start and end step
2. Is the source node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING)
3. Is the destination node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING)
4. Are there any Access Control List rules in place that prevent the application of this IER
5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level)
For red agent IERs, the application of IERs between a source and destination follows a number of subtly different rules. Specifically:
1. Does the current simulation time step fall between IER start and end step
2. Is the source node operational, and is the service (protocol / port) associated with the IER (a) present on that node and (b) already in a compromised state
3. Is the destination node operational, and is the service (protocol / port) associated with the IER present on that node
4. Are there any Access Control List rules in place that prevent the application of this IER
5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level)
Assuming the rules pass, the IER is applied to all relevant links (based on use of OSPF) between source and destination.
Node Pattern-of-Life
********************
Every node can be impacted (i.e. have a status change applied to it) by either green agent pattern-of-life or red agent pattern-of-life. This is distinct from IERs, and allows for attacks (and defence) to be modelled purely within the confines of a node.
The status changes that can be made to a node are as follows:
* All Nodes:
* Operating State:
* ON
* OFF
* RESETTING - when a status of resetting is entered, the node will automatically exit this state after a number of steps (as defined by the nodeResetDuration configuration item) after which it returns to an ON state
* Active Nodes and Service Nodes:
* Operating System State:
* GOOD
* PATCHING - when a status of patching is entered, the node will automatically exit this state after a number of steps (as defined by the osPatchingDuration configuration item) after which it returns to a GOOD state
* COMPROMISED
* Service Nodes only:
* Service State (for any associated service):
* GOOD
* PATCHING - when a status of patching is entered, the service will automatically exit this state after a number of steps (as defined by the servicePatchingDuration configuration item) after which it returns to a GOOD state
* COMPROMISED
* OVERWHELMED
Access Control List modelling
*****************************
An Access Control List (ACL) is modelled to provide the means to manage traffic flows in the system. This will allow defensive agents the means to turn on / off rules, or potentially create new rules, to counter an attack.
The ACL follows a standard network firewall format. For example:
.. list-table:: ACL example
:widths: 25 25 25 25 25
:header-rows: 1
* - Permission
- Source IP
- Dest IP
- Protocol
- Port
* - DENY
- 192.168.1.2
- 192.168.1.3
- HTTPS
- 443
* - ALLOW
- 192.168.1.4
- ANY
- SMTP
- 25
* - DENY
- ANY
- 192.168.1.5
- ANY
- ANY
All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or ALLOW for the same parameters will override an earlier entry.
Observation Spaces
******************
The OpenAI Gym observation space provides the status of all nodes and links across the whole system:
* Nodes (in terms of operating state, operating system state, and services state)
* Links (in terms of current loading for each service/protocol)
An example observation space is provided below:
.. list-table:: Observation Space example
:widths: 25 25 25 25 25 25
:header-rows: 1
* -
- ID
- Operating State
- O/S State
- Service / Protocol A
- Service / Protocol B
* - Node A
- 1
- 1
- 1
- 1
- 1
* - Node B
- 2
- 1
- 3
- 1
- 1
* - Node C
- 3
- 2
- 1
- 3
- 2
* - Link 1
- 5
- 1
- 1
- 0
- 10000
* - Link 2
- 6
- 1
- 1
- 0
- 10000
* - Link 3
- 7
- 1
- 1
- 0
- 0
The observation space is a 6 x 5 Box type (OpenAI Gym Space) in this example. This is made up from the node and link information detailed below.
For the nodes, the following values are represented:
* ID
* Operating State:
* 1 = ON
* 2 = OFF
* 3 = RESETTING
* O/S State:
* 1 = GOOD
* 2 = PATCHING
* 3 = COMPROMISED
* Service State:
* 1 = GOOD
* 2 = PATCHING
* 3 = COMPROMISED
* 4 = OVERWHELMED
(Note that each service available in the network is provided as a column, although not all nodes may utilise all services)
For the links, the following statuses are represented:
* ID
* Operating State = N/A
* O/S State = N/A
* Protocol = loading in bits/s
Action Spaces
**************
The action space available to the blue agent comes in two types:
1. Node-based
2. Access Control List
The choice of action space used during a training session is determined in the config_[name].yaml file.
**Node-Based**
The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym multidiscrete type, as follows:
* [0, num nodes] - Node ID (0 = nothing, node ID)
* [0, 3] - What property it's acting on (0 = nothing, 1 = state, 2 = O/S state, 3 = service state)
* [0, 3] - Action on property (0 = nothing, 1 = on, 2 = off, 3 = reset / patch)
* [0, num services] - Resolves to service ID (0 = nothing, resolves to service)
**Access Control List**
The blue agent is able to influence the configuration of the Access Control List rule set (which implements a system-wide firewall). In this instance, the action space is an OpenAI multidiscrete type, as follows:
* [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
* [0, 1] - Permission (0 = DENY, 1 = ALLOW)
* [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses)
* [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses)
* [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol)
* [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
Rewards
*******
A reward value is presented back to the blue agent on the conclusion of every step. The reward value is calculated via two methods which combine to give the total value:
1. Node and service status
2. IER status
**Node and service status**
On every step, the status of each node is compared against both a reference environment (simulating the situation if the red and blue agents had not impacted the environment)
and the before and after state of the environment. If the comparison against the reference environment shows no difference, then the score provided is "AllOK". If there is a
difference with respect to the reference environment, the before and after states are compared, and a score determined. See :ref:`config` for details of reward values.
**IER status**
On every step, the full IER set is examined to determine whether green and red agent IERs are being permitted to run. Any red agent IERs running incur a penalty; any green agent
IERs not permitted to run also incur a penalty. See :ref:`config` for details of reward values.
Future Enhancements
*******************
The PrimAITE project has an ambition to include the following enhancements in future releases:
* Integration with a suitable standardised framework to allow multi-agent integration
* Integration with external threat emulation tools, either using off-line data, or integrating at runtime
* Provision of data such that agents can construct alternative observation spaces (as an alternative to the default PrimAITE observation space)
* Introduction of a testing phase (post training) to evaluate the effectiveness of the training

View File

@@ -0,0 +1,28 @@
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html
# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
project = 'PrimAITE'
copyright = '2022, jashort'
author = 'jashort'
release = '0.1.0'
# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
extensions = ['sphinx_rtd_theme']
templates_path = ['_templates']
exclude_patterns = []
# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
html_theme = 'sphinx_rtd_theme'
html_static_path = ['_static']

View File

@@ -0,0 +1,261 @@
.. _config:
The Config Files Explained
==========================
PrimAITE uses two configuration files for its operation:
* config_main.yaml - used to define the top-level settings of the PrimAITE environment, and the training session that is to be run.
* config_[name].yaml - used to define the low-level settings of a training session, including the network laydown, green / red agent information exchange requirements (IERSs), Access Control Rules, Action Space type, and the number of steps in each episode.
config_main.yaml:
*****************
The config_main.yaml file consists of the following attributes:
**Generic Config Values**
* **agentIdentifier** [enum]
This identifies the agent to use for the training session. Select from one of the following:
* GENERIC - Where a user developed agent is to be used
* STABLE_BASELINES3_PPO - Use a SB3 PPO agent
* STABLE_BASELINES3_A2C - use a SB3 A2C agent
* **numEpisodes** [int]
This defines the number of episodes that the agent will train over. Each episode consists of a number of steps (with step number defined in the config_[name].yaml file)
* **timeDelay** [int]
The time delay (in milliseconds) to take between each step when training a GENERIC agent
* **configFilename** [filename]
The name of the config_[name].yaml file to use for this training session
* **observationSpaceHighValue** [int]
The high value to use for values in the observation space. This is set to 1000000000 by default, and should not need changing in most cases
**Reward-Based Config Values**
* **Generic [allOk]** [int]
The score to give when the current situation (for a given component) is no different from that expected in the baseline (i.e. as though no blue or red agent actions had been undertaken)
* **Node Operating State [offShouldBeOn]** [int]
The score to give when the node should be on, but is off
* **Node Operating State [offShouldBeResetting]** [int]
The score to give when the node should be resetting, but is off
* **Node Operating State [onShouldBeOff]** [int]
The score to give when the node should be off, but is on
* **Node Operating State [onShouldBeResetting]** [int]
The score to give when the node should be resetting, but is on
* **Node Operating State [resettingShouldBeOn]** [int]
The score to give when the node should be on, but is resetting
* **Node Operating State [resettingShouldBeOff]** [int]
The score to give when the node should be off, but is resetting
* **Node Operating System or Service State [goodShouldBePatching]** [int]
The score to give when the state should be patching, but is good
* **Node Operating System or Service State [goodShouldBeCompromised]** [int]
The score to give when the state should be compromised, but is good
* **Node Operating System or Service State [goodShouldBeOverwhelmed]** [int]
The score to give when the state should be overwhelmed, but is good
* **Node Operating System or Service State [patchingShouldBeGood]** [int]
The score to give when the state should be good, but is patching
* **Node Operating System or Service State [patchingShouldBeCompromised]** [int]
The score to give when the state should be compromised, but is patching
* **Node Operating System or Service State [patchingShouldBeOverwhelmed]** [int]
The score to give when the state should be overwhelmed, but is patching
* **Node Operating System or Service State [compromisedShouldBeGood]** [int]
The score to give when the state should be good, but is compromised
* **Node Operating System or Service State [compromisedShouldBePatching]** [int]
The score to give when the state should be patching, but is compromised
* **Node Operating System or Service State [compromisedShouldBeOverwhelmed]** [int]
The score to give when the state should be overwhelmed, but is compromised
* **Node Operating System or Service State [compromised]** [int]
The score to give when the state is compromised
* **Node Operating System or Service State [overwhelmedShouldBeGood]** [int]
The score to give when the state should be good, but is overwhelmed
* **Node Operating System or Service State [overwhelmedShouldBePatching]** [int]
The score to give when the state should be patching, but is overwhelmed
* **Node Operating System or Service State [overwhelmedShouldBeCompromised]** [int]
The score to give when the state should be compromised, but is overwhelmed
* **Node Operating System or Service State [overwhelmed]** [int]
The score to give when the state is overwhelmed
* **IER Status [redIerRunning]** [int]
The score to give when a red agent IER is permitted to run
* **IER Status [greenIerBlocked]** [int]
The score to give when a green agent IER is prevented from running
**Patching / Reset Durations**
* **osPatchingDuration** [int]
The number of steps to take when patching an Operating System
* **nodeResetDuration** [int]
The number of steps to take when resetting a node's operating state
* **servicePatchingDuration** [int]
The number of steps to take when patching a service
config_[name].yaml:
*******************
The config_[name].yaml file consists of the following attributes:
* **itemType: ACTIONS** [enum]
Determines whether a NODE or ACL action space format is adopted for the training session
* **itemType: STEPS** [int]
Determines the number of steps to run in each episode of the training session
* **itemType: PORTS** [int]
Provides a list of ports modelled in this training session
* **itemType: SERVICES** [freetext]
Provides a list of services modelled in this training session
* **itemType: NODE**
Defines a node included in the system laydown being simulated. It should consist of the following attributes:
* **id** [int]: Unique ID for this YAML item
* **name** [freetext]: Human-readable name of the component
* **baseType** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services
* **nodeType** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER
* **priority** [enum]: Provides a priority for each node. Can be one of P1, P2, P3, P4 or P5 (which P1 being the highest)
* **hardwareState** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING
* **ipAddress** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx
* **softwareState** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED
* **services**: For each service associated with the node:
* **name** [freetext]: Free-text name of the service, but must match one of the services defined for the system in the services list
* **port** [int]: Integer value of the port related to this service, but must match one of the ports defined for the system in the ports list
* **state** [enum]: The initial state of the service. Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED
* **itemType: LINK**
Defines a link included in the system laydown being simulated. It should consist of the following attributes:
* **id** [int]: Unique ID for this YAML item
* **name** [freetext]: Human-readable name of the component
* **bandwidth** [int]: The bandwidth (in bits/s) of the link
* **source** [int]: The ID of the source node
* **destination** [int]: The ID of the destination node
* **itemType: GREEN_IER**
Defines a green agent Information Exchange Requirement (IER). It should consist of:
* **id** [int]: Unique ID for this YAML item
* **startStep** [int]: The start step (in the episode) for this IER to begin
* **endStep** [int]: The end step (in the episode) for this IER to finish
* **load** [int]: The load (in bits/s) for this IER to apply to links
* **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list
* **port** [int]: The port that the protocol is running on. This must match a value in the ports list
* **source** [int]: The ID of the source node
* **destination** [int]: The ID of the destination node
* **missionCriticality** [enum]: The mission criticality of this IER (with 5 being highest, 1 lowest)
* **itemType: RED_IER**
Defines a red agent Information Exchange Requirement (IER). It should consist of:
* **id** [int]: Unique ID for this YAML item
* **startStep** [int]: The start step (in the episode) for this IER to begin
* **endStep** [int]: The end step (in the episode) for this IER to finish
* **load** [int]: The load (in bits/s) for this IER to apply to links
* **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list
* **port** [int]: The port that the protocol is running on. This must match a value in the ports list
* **source** [int]: The ID of the source node
* **destination** [int]: The ID of the destination node
* **missionCriticality** [enum]: Not currently used. Default to 0
* **itemType: GREEN_POL**
Defines a green agent pattern-of-life instruction. It should consist of:
* **id** [int]: Unique ID for this YAML item
* **startStep** [int]: The start step (in the episode) for this PoL to begin
* **endStep** [int]: Not currently used. Default to same as start step
* **node** [int]: The ID of the node to apply the PoL to
* **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE
* **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list
* **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state)
* **itemType: RED_POL**
Defines a red agent pattern-of-life instruction. It should consist of:
* **id** [int]: Unique ID for this YAML item
* **startStep** [int]: The start step (in the episode) for this PoL to begin
* **endStep** [int]: Not currently used. Default to same as start step
* **node** [int]: The ID of the node to apply the PoL to
* **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE
* **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list
* **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state)
* **isEntryNode** [bool]: Defines whether the node affected is an entry node to the system
* **itemType: ACL_RULE**
Defines an initial Access Control List (ACL) rule. It should consist of:
* **id** [int]: Unique ID for this YAML item
* **permission** [enum]: Defines either an allow or deny rule. Value must be either DENY or ALLOW
* **source** [IP address]: Defines the source IP address for the rule in xxx.xxx.xxx.xxx format
* **destination** [IP address]: Defines the destination IP address for the rule in xxx.xxx.xxx.xxx format
* **protocol** [freetext]: Defines the protocol for the rule. Must match a value in the services list
* **port** [int]: Defines the port for the rule. Must match a value in the ports list

View File

@@ -0,0 +1,26 @@
.. _dependencies:
PrimAITE Dependencies
=====================
PrimAITE is built with the following versions of dependencies:
* Python 3.10.9
* PyYAML 6.0
* numpy 1.23.5
* networkx 2.8.8
* gym 0.21.0
* matplotlib 3.6.2
* stable_baselines_3 1.6.2
The latest release of PrimAITE has been tested against the following versions of dependencies:
* Python 3.10.9
* PyYAML 6.0
* numpy 1.23.5
* networkx 2.8.8
* gym 0.21.0
* matplotlib 3.6.2
* stable_baselines_3 1.6.2

View File

@@ -0,0 +1,42 @@
.. PrimAITE documentation master file, created by
sphinx-quickstart on Thu Dec 8 09:51:18 2022.
You can adapt this file completely to your liking, but it should at least
contain the root `toctree` directive.
Welcome to PrimAITE's documentation
====================================
What is PrimAITE?
------------------------
PrimAITE (Primary-level AI Training Environment) is a simulation environment for training AI under the ARCD programme. It incorporates the functionality required of a Primary-level environment, as specified in the Dstl ARCD Training Environment Matrix document:
* The ability to model a relevant platform / system context;
* The ability to model key characteristics of a platform / system by representing connections, IP addresses, ports, traffic loading, operating systems, services and processes;
* Operates at machine-speed to enable fast training cycles.
PrimAITE aims to evolve into an ARCD environment that could be used as the follow-on from Reception level approaches (e.g. YAWNING TITAN), and help bridge the Sim-to-Real gap into Secondary level environments (e.g. IMAGINARY YAK).
This is similar to the approach taken by FVEY international partners (e.g. AUS CyBORG, US NSA FARLAND and CAN CyGil). These environments are referenced by the Dstl ARCD Agent Training Environments Knowledge Transfer document (TR141342).
What is PrimAITE built with
--------------------------------------
* `OpenAI's Gym <https://gym.openai.com/>`_ is used as the basis for AI blue agent interaction with the PrimAITE environment
* `Networkx <https://github.com/networkx/networkx>`_ is used as the underlying data structure used for the PrimAITE environment
* `Stable Baselines 3 <https://github.com/DLR-RM/stable-baselines3>`_ is used as a default source of RL algorithms (although PrimAITE is not limited to SB3 agents)
Where next?
------------
The best place to start is :ref:`about`
.. toctree::
:maxdepth: 8
:caption: Contents:
about
dependencies
config
training
results

View File

@@ -0,0 +1,42 @@
.. _results:
Results, Output and Logging from PrimAITE
=========================================
PrimAITE produces four types of data:
* Outputs - Results
* Outputs - Diagrams
* Outputs - Saved agents
* Logging
Outputs can be found in the *[Install Directory]\\PRIMAITE\\PRIMAITE\\outputs* directory
Logging can be found in the *[Install Directory]\\PRIMAITE\\PRIMAITE\\logs* directory
**Outputs - Results**
PrimAITE automatically creates two sets of results from each training session, and stores them in the *Results* folder:
* Average reward per episode - a csv file listing the average reward for each episode of the training session. This provides an indication of the change, over a training session, of the reward value
* All transactions - a csv file listing the following values for every step of every episode:
* Timestamp
* Episode number
* Step number
* Initial observation space (before red and blue agent actions have been taken). Individual elements of the observation space are presented in the format OSI_X_Y
* Resulting observation space (after the red and blue agent actions have been taken) Individual elements of the observation space are presented in the format OSN_X_Y
* Reward value
* Action space (as presented by the blue agent on this step). Individual elements of the action space are presented in the format AS_X
**Outputs - Diagrams**
For each training run, PrimAITE automatically creates a visual of the system / network laydown configuration, and stores it in the *Diagrams* folder.
**Outputs - Saved agents**
For each training run, assuming the agent being trained implements the *save()* function and this function is called by the code, PrimAITE automatically saves the agent state and stores it in the *agents* folder.
**Logging**
PrimAITE also provides output logs (for diagnosis) using the Python Logging package. These can be found in the *[Install Directory]\\PRIMAITE\\PRIMAITE\\logs* directory

View File

@@ -0,0 +1,88 @@
.. _training:
Running a PrimAITE Training Session
===================================
A PrimAITE training session will usually be associated with a "Training Use Case Profile". This document will present:
* The Use Case name, default number of steps in a training episode and default number of episodes in a training session. The number of steps and episodes can be modified in the configuration files
* The system laydown being modelled
* The objectives of the session (steady-state), the red agent and the blue agent (in a defensive role)
* The green agent pattern-of-life profile
* The red agent attack profile
* The observation space definition
* The action space definition
* Agent integration guidance
* Initial Access Control List settings (if applicable)
* The reward function definition
**Integrating a user defined blue agent**
Integrating a blue agent with PrimAITE requires some modification of the code within the main.py file. The main.py file consists of a number of functions, each of which will invoke training for a particular agent. These are:
* Generic (run_generic)
* Stable Baselines 3 PPO (run_stable_baselines3_ppo)
* Stable Baselines 3 A2C (run_stable_baselines3_a2c)
The selection of which agent type to use is made via the config_main.yaml file. In order to train a user generated agent,
the run_generic function should be selected, and should be modified (typically) to be:
.. code:: python
agent = MyAgent(environment, max_steps)
for episode in range(0, num_episodes):
agent.learn()
env.close()
save_agent(agent)
Where:
* *MyAgent* is the user created agent
* *environment* is the PrimAITE environment
* *max_steps* is the number of steps in an episode, as defined in the config_[name].yaml file
* *num_episodes* is the number of episodes in the training session, as defined in the config_main.yaml file
* the *.learn()* function should be defined in the user created agent
* the *env.close()* function is defined within PrimAITE
* the *save_agent()* assumes that a *save()* function has been defined in the user created agent. If not, this line can be ommitted (although it is encouraged, since it will allow the agent to be saved and ported)
The code below provides a suggested format for the learn() function within the user created agent.
It's important to include the *self.environment.reset()* call within the episode loop in order that the
environment is reset between episodes. Note that the example below should not be considered exhaustive.
.. code:: python
def learn(self) :
# pre-reqs
# reset the environment
self.environment.reset()
done = False
for step in range(max_steps):
# calculate the action
action = ...
# execute the environment step
new_state, reward, done, info = self.environment.step(action)
# algorithm updates
...
# update to our new state
state = new_state
# if done, finish episode
if done == True:
break
**Running the training session**
In order to execute a training session, carry out the following steps:
1. Navigate to "[Install directory]\\PRIMAITE\\PRIMAITE\\”
2. Start a console window (type “CMD” in path window, or start a console window first and navigate to “[Install Directory]\\PRIMAITE\\PRIMAITE\\”)
3. Type “python main.py”
4. Training will start with an output indicating the current episode, and average reward value for the episode

View File

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

View File

@@ -0,0 +1,989 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Main environment module containing the PRIMmary AI Training Evironment (PRIMAITE) class
"""
import numpy as np
import networkx as nx
import copy
import csv
import yaml
import os.path
import logging
from gym import Env, spaces
from matplotlib import pyplot as plt
from datetime import datetime
from common.enums import *
from links.link import Link
from pol.ier import IER
from nodes.node_state_instruction import NodeStateInstruction
from pol.green_pol import apply_iers, apply_node_pol
from pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
from nodes.active_node import ActiveNode
from nodes.passive_node import PassiveNode
from nodes.service_node import ServiceNode
from common.service import Service
from acl.access_control_list import AccessControlList
from environment.reward import calculate_reward_function
from transactions.transaction import Transaction
class PRIMAITE(Env):
"""
PRIMmary AI Training Evironment (PRIMAITE) class
"""
# Observation / Action Space contants
OBSERVATION_SPACE_FIXED_PARAMETERS = 3
ACTION_SPACE_NODE_PROPERTY_VALUES = 4
ACTION_SPACE_NODE_ACTION_VALUES = 4
ACTION_SPACE_ACL_ACTION_VALUES = 3
ACTION_SPACE_ACL_PERMISSION_VALUES = 2
OBSERVATION_SPACE_HIGH_VALUE = 1000000 # Highest value within an observation space
def __init__(self, _config_values, _transaction_list):
"""
Init
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
"""
super(PRIMAITE, self).__init__()
# Take a copy of the config values
self.config_values = _config_values
# Number of steps in an episode
self.episode_steps = 0
# Transaction list
self.transaction_list = _transaction_list
# The agent in use
self.agent_identifier = self.config_values.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes = {}
# Create a dictionary to hold a reference set of nodes
self.nodes_reference = {}
# Create a dictionary to hold all the links
self.links = {}
# Create a dictionary to hold a reference set of links
self.links_reference = {}
# Create a dictionary to hold all the green IERs (this will come from an external source)
self.green_iers = {}
# Create a dictionary to hold all the node PoLs (this will come from an external source)
self.node_pol = {}
# Create a dictionary to hold all the red agent IERs (this will come from an external source)
self.red_iers = {}
# Create a dictionary to hold all the red agent node PoLs (this will come from an external source)
self.red_node_pol = {}
# Create the Access Control List
self.acl = AccessControlList()
# Create a list of services (enums)
self.services_list = []
# Create a list of ports
self.ports_list = []
# Create graph (network)
self.network = nx.MultiGraph()
# Create a graph (network) reference
self.network_reference = nx.MultiGraph()
# Create step count
self.step_count = 0
# Create step info dictionary
self.step_info = {}
# Total reward
self.total_reward = 0
# Average reward
self.average_reward = 0
# Episode count
self.episode_count = 0
# Number of nodes - gets a value by examining the nodes dictionary after it's been populated
self.num_nodes = 0
# Number of links - gets a value by examining the links dictionary after it's been populated
self.num_links = 0
# Number of services - gets a value when config is loaded
self.num_services = 0
# Number of ports - gets a value when config is loaded
self.num_ports = 0
# The action type
self.action_type = 0
# Open the config file and build the environment laydown
try:
self.config_file = open("config/" + self.config_values.config_filename_use_case, "r")
self.config_data = yaml.safe_load(self.config_file)
self.load_config()
except Exception as e:
logging.error("Could not load the environment configuration")
logging.error("Exception occured", exc_info=True)
# Store the node objects as node attributes
# (This is so we can access them as objects)
for node in self.network:
self.network.nodes[node]["self"] = node
for node in self.network_reference:
self.network_reference.nodes[node]["self"] = node
self.num_nodes = len(self.nodes)
self.num_links = len(self.links)
# Visualise in PNG
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")
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")
plt.clf()
except Exception as a:
logging.error("Could not save network diagram")
logging.error("Exception occured", exc_info=True)
print("Could not save network diagram")
# Define Observation Space
# x = number of nodes and links (i.e. items)
# y = number of parameters to be sent
# For each item, we send:
# - [For Nodes] | [For Links]
# - node ID | link ID
# - operating state | N/A
# - operating system state | N/A
# - service A state | service A loading
# - service B state | service B loading
# - service C state | service C loading
# - service D state | service D loading
# - service E state | service E loading
# - service F state | service F loading
# - service G state | service G loading
# Calculate the number of items that need to be included in the observation space
num_items = self.num_links + self.num_nodes
# Set the number of observation parameters, being # of services plus id, operating sytem system and O/S state (i.e. 3)
self.num_observation_parameters = self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
# Define the observation shape
self.observation_shape = (num_items, self.num_observation_parameters)
self.observation_space = spaces.Box(low=0,
high=self.config_values.observation_space_high_value,
shape=self.observation_shape,
dtype=np.int64)
# This is the observation that is sent back via the rest and step functions
self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ACTION_TYPE.NODE:
logging.info("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
# [0, 3] - what property it's acting on (0 = nothing, state, o/s state, service state)
# [0, 3] - action on property (0 = nothing, On, Off, Reset / Patch)
# [0, num services] - resolves to service ID (0 = nothing, resolves to service)
self.action_space = spaces.MultiDiscrete([self.num_nodes, self.ACTION_SPACE_NODE_PROPERTY_VALUES, self.ACTION_SPACE_NODE_ACTION_VALUES, self.num_services])
else:
logging.info("Action space type ACL selected")
# Terms (for ACL action space):
# [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
# [0, 1] - Permission (0 = DENY, 1 = ALLOW)
# [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses)
# [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses)
# [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol)
# [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
self.action_space = spaces.MultiDiscrete([self.ACTION_SPACE_ACL_ACTION_VALUES, self.ACTION_SPACE_ACL_PERMISSION_VALUES, self.num_nodes + 1, self.num_nodes + 1, self.num_services + 1, self.num_ports + 1])
# 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='')
self.csv_writer = csv.writer(self.csv_file)
self.csv_writer.writerow(header)
except Exception as e:
logging.error("Could not create csv file to hold average reward per episode")
logging.error("Exception occured", exc_info=True)
def reset(self):
"""
AI Gym Reset function
Returns:
Environment observation space (reset)
"""
csv_data = self.episode_count, self.average_reward
self.csv_writer.writerow(csv_data)
self.episode_count += 1
# Don't need to reset links, as they are cleared and recalculated every step
# Clear the ACL
self.init_acl()
# Reset the node statuses and recreate the ACL from config
# Does this for both live and reference nodes
self.reset_environment()
# Reset counters and totals
self.total_reward = 0
self.step_count = 0
self.average_reward = 0
# Update observations space and return
self.update_environent_obs()
return self.env_obs
def step(self, action):
"""
AI Gym Step function
Args:
action: Action space from agent
Returns:
env_obs: Observation space
reward: Reward value for this step
done: Indicates episode is complete if True
step_info: Additional information relating to this step
"""
if self.step_count == 0:
print("Episode: " + str(self.episode_count) + " running")
# TEMP
done = False
self.step_count += 1
#print("Episode step: " + str(self.stepCount))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
link_value.clear_traffic()
# Create a Transaction (metric) object for this step
transaction = Transaction(datetime.now(), self.agent_identifier, self.episode_count, self.step_count)
# Load the initial observation space into the transaction
transaction.set_obs_space_pre(copy.deepcopy(self.env_obs))
# Load the action space into the transaction
transaction.set_action_space(copy.deepcopy(action))
# 1. Perform any time-based activities (e.g. a component moving from patching to good)
self.apply_time_based_updates()
# 2. Apply PoL
apply_node_pol(self.nodes, self.node_pol, self.step_count) # Node PoL
apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count) # Network PoL
# Take snapshots of nodes and links
self.nodes_post_pol = copy.deepcopy(self.nodes)
self.links_post_pol = copy.deepcopy(self.links)
# Reference
apply_node_pol(self.nodes_reference, self.node_pol, self.step_count) # Node PoL
apply_iers(self.network_reference, self.nodes_reference, self.links_reference, self.green_iers, self.acl, self.step_count) # Network PoL
# 3. Implement Red Action
apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
apply_red_agent_node_pol(self.nodes, self.red_iers, self.red_node_pol, self.step_count)
# Take snapshots of nodes and links
self.nodes_post_red = copy.deepcopy(self.nodes)
self.links_post_red = copy.deepcopy(self.links)
# 4. Implement Blue Action
self.interpret_action_and_apply(action)
# 5. Reapply normal and Red agent IER PoL, as we need to see what effect the blue agent action has had (if any) on link status
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
link_value.clear_traffic()
apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count)
apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
# Take snapshots of nodes and links
self.nodes_post_blue = copy.deepcopy(self.nodes)
self.links_post_blue = copy.deepcopy(self.links)
# 6. Calculate reward signal (for RL)
reward = calculate_reward_function(self.nodes_post_pol, self.nodes_post_blue, self.nodes_reference, self.green_iers, self.red_iers, self.step_count, self.config_values)
#print("Step reward: " + str(reward))
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
print("Average reward: " + str(self.average_reward))
# Load the reward into the transaction
transaction.set_reward(reward)
# 7. Output Verbose
#self.output_link_status()
# 8. Update env_obs
self.update_environent_obs()
# Load the new observation space into the transaction
transaction.set_obs_space_post(copy.deepcopy(self.env_obs))
# 9. Add the transaction to the list of transactions
self.transaction_list.append(copy.deepcopy(transaction))
# Return
return self.env_obs, reward, done, self.step_info
def __close__(self):
"""
Override close function
"""
self.csv_file.close()
self.config_file.close()
def init_acl(self):
"""
Initialise the Access Control List
"""
self.acl.remove_all_rules()
def output_link_status(self):
"""
Output the link status of all links to the console
"""
for link_key, link_value in self.links.items():
print("Link ID: " + link_value.get_id())
for protocol in link_value.get_protocol_list():
print(" Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load()))
def interpret_action_and_apply(self, _action):
"""
Applies agent actions to the nodes and Access Control List
Args:
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.action_type == ACTION_TYPE.NODE:
self.apply_actions_to_nodes(_action)
else:
self.apply_actions_to_acl(_action)
def apply_actions_to_nodes(self, _action):
"""
Applies agent actions to the nodes
Args:
_action: The action space from the agent
"""
node_id = _action[0]
node_property = _action[1]
property_action = _action[2]
service_index = _action[3]
# Check that the action is requesting a valid node
try:
node = self.nodes[str(node_id)]
except:
return
if node_property == 0:
# This is the do nothing action
return
elif node_property == 1:
# This is an action on the node Operating State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Turn on (only applicable if it's OFF, not if it's patching)
if node.get_state() == HARDWARE_STATE.OFF:
node.turn_on()
elif property_action == 2:
# Turn off
node.turn_off()
elif property_action == 3:
# Reset (only applicable if it's ON)
if node.get_state() == HARDWARE_STATE.ON:
node.reset()
else:
return
elif node_property == 2:
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
# This is an action on the node Operating System State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_os_state(SOFTWARE_STATE.PATCHING)
else:
# Node is not of Active or Service Type
return
elif node_property == 3:
# This is an action on a node Service State
if isinstance(node, ServiceNode):
# This is an action on a node Service State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_service_state(self.services_list[service_index], SOFTWARE_STATE.PATCHING)
else:
# Node is not of Service Type
return
else:
return
def apply_actions_to_acl(self, _action):
"""
Applies agent actions to the Access Control List [TO DO]
Args:
_action: The action space from the agent
"""
action_decision = _action[0]
action_permission = _action[1]
action_source_ip = _action[2]
action_destination_ip = _action[3]
action_protocol = _action[4]
action_port = _action[5]
if action_decision == 0:
# It's decided to do nothing
return
else:
# It's decided to create a new ACL rule or remove an existing rule
# Permission value
if action_permission == 0:
acl_rule_permission = "DENY"
else:
acl_rule_permission = "ALLOW"
# Source IP value
if action_source_ip == 0:
acl_rule_source = "ANY"
else:
node = list(self.nodes.values())[action_source_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_source = node.get_ip_address()
else:
return
# Destination IP value
if action_destination_ip == 0:
acl_rule_destination = "ANY"
else:
node = list(self.nodes.values())[action_destination_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_destination = node.get_ip_address()
else:
return
# Protocol value
if action_protocol == 0:
acl_rule_protocol = "ANY"
else:
acl_rule_protocol = self.services_list[action_protocol - 1]
# Port value
if action_port == 0:
acl_rule_port = "ANY"
else:
acl_rule_port = self.ports_list[action_port - 1]
# Now add or remove
if action_decision == 1:
# Add the rule
self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
elif action_decision == 2:
# Remove the rule
self.acl.remove_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
else:
return
def apply_time_based_updates(self):
"""
Updates anything that needs to count down and then change state (e.g. reset / patching status)
"""
for node_key, node in self.nodes.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
node.update_resetting_status()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
node.update_os_patching_status()
else:
pass
else:
pass
if isinstance(node, ServiceNode):
node.update_services_patching_status()
else:
pass
for node_key, node in self.nodes_reference.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
node.update_resetting_status()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
node.update_os_patching_status()
else:
pass
else:
pass
if isinstance(node, ServiceNode):
node.update_services_patching_status()
else:
pass
def update_environent_obs(self):
"""
# Updates the observation space based on the node and link status
"""
item_index = 0
# Do nodes first
for node_key, node in self.nodes.items():
self.env_obs[item_index][0] = int(node.get_id())
self.env_obs[item_index][1] = node.get_state().value
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
self.env_obs[item_index][2] = node.get_os_state().value
else:
self.env_obs[item_index][2] = 0
service_index = 3
if isinstance(node, ServiceNode):
for service in self.services_list:
if node.has_service(service):
self.env_obs[item_index][service_index] = node.get_service_state(service).value
else:
self.env_obs[item_index][service_index] = 0
service_index += 1
else:
# Not a service node
for service in self.services_list:
self.env_obs[item_index][service_index] = 0
service_index += 1
item_index += 1
# Now do links
for link_key, link in self.links.items():
self.env_obs[item_index][0] = int(link.get_id())
self.env_obs[item_index][1] = 0
self.env_obs[item_index][2] = 0
protocol_list = link.get_protocol_list()
protocol_index = 0
for protocol in protocol_list:
self.env_obs[item_index][protocol_index + 3] = protocol.get_load()
protocol_index += 1
item_index += 1
def load_config(self):
"""
# Loads config data in order to build the environment configuration
"""
for item in self.config_data:
if item["itemType"] == "NODE":
# Create a node
self.create_node(item)
elif item["itemType"] == "LINK":
# Create a link
self.create_link(item)
elif item["itemType"] == "GREEN_IER":
# Create a Green IER
self.create_green_ier(item)
elif item["itemType"] == "GREEN_POL":
# Create a Green PoL
self.create_green_pol(item)
elif item["itemType"] == "RED_IER":
# Create a Red IER
self.create_red_ier(item)
elif item["itemType"] == "RED_POL":
# Create a Red PoL
self.create_red_pol(item)
elif item["itemType"] == "ACL_RULE":
# Create an ACL rule
self.create_acl_rule(item)
elif item["itemType"] == "SERVICES":
# Create the list of services
self.create_services_list(item)
elif item["itemType"] == "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)
pass
logging.info("Environment configuration loaded")
print("Environment configuration loaded")
def create_node(self, item):
"""
Creates a node from config data
Args:
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_name = item["name"]
node_base_type = item["baseType"]
node_type = TYPE[item["nodeType"]]
node_priority = PRIORITY[item["priority"]]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
if node_base_type == "PASSIVE":
node = PassiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, self.config_values)
elif node_base_type == "ACTIVE":
# Active nodes have IP address and operating system state
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node = ActiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, self.config_values)
elif node_base_type == "SERVICE":
# Service nodes have IP address, operating system state and list of services
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node = ServiceNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, self.config_values)
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
service_port = service["port"]
service_state = SOFTWARE_STATE[service["state"]]
node.add_service(Service(service_protocol, service_port, service_state))
else:
# Bad formatting
pass
# Copy the node for the reference version
node_ref = copy.deepcopy(node)
# Add node to node dictionary
self.nodes[node_id] = node
# Add reference node to reference node dictionary
self.nodes_reference[node_id] = node_ref
# Add node to network
self.network.add_nodes_from([node])
# Add node to network (reference)
self.network_reference.add_nodes_from([node_ref])
def create_link(self, item):
"""
Creates a link from config data
Args:
item: A config data item
"""
link_id = item["id"]
link_name = item["name"]
link_bandwidth = item["bandwidth"]
link_source = item["source"]
link_destination = item["destination"]
source_node = self.nodes[link_source]
dest_node = self.nodes[link_destination]
# Add link to network
self.network.add_edge(source_node, dest_node, id=link_name)
# Add link to link dictionary
self.links[link_name] = Link(link_id, link_bandwidth, source_node.get_name(), dest_node.get_name(), self.services_list)
# Reference
source_node_ref = self.nodes_reference[link_source]
dest_node_ref = self.nodes_reference[link_destination]
# Add link to network (reference)
self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name)
# Add link to link dictionary (reference)
self.links_reference[link_name] = Link(link_id, link_bandwidth, source_node_ref.get_name(), dest_node_ref.get_name(), self.services_list)
def create_green_ier(self, item):
"""
Creates a green IER from config data
Args:
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
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"]
# Create IER and add to green IER dictionary
self.green_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
def create_red_ier(self, item):
"""
Creates a red IER from config data
Args:
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
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"]
# Create IER and add to red IER dictionary
self.red_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
def create_green_pol(self, item):
"""
Creates a green PoL object from config data
Args:
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_node = item["node"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
else:
pol_state = SOFTWARE_STATE[item["state"]]
self.node_pol[pol_id] = NodeStateInstruction(pol_id, pol_start_step, pol_end_step, pol_node, pol_type, pol_protocol, pol_state)
def create_red_pol(self, item):
"""
Creates a red PoL object from config data
Args:
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_node = item["node"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
else:
pol_state = SOFTWARE_STATE[item["state"]]
pol_is_entry_node = item["isEntryNode"]
self.red_node_pol[pol_id] = NodeStateInstruction(pol_id, pol_start_step, pol_end_step, pol_node, pol_type, pol_protocol, pol_state, pol_is_entry_node)
def create_acl_rule(self, item):
"""
Creates an ACL rule from config data
Args:
item: A config data item
"""
acl_rule_permission = item["permission"]
acl_rule_source = item["source"]
acl_rule_destination = item["destination"]
acl_rule_protocol = item["protocol"]
acl_rule_port = item["port"]
self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
def create_services_list(self, services):
"""
Creates a list of services (enum) from config data
Args:
item: A config data item representing the services
"""
service_list = services["serviceList"]
for service in service_list:
service_name = service["name"]
self.services_list.append(service_name)
# Set the number of services
self.num_services = len(self.services_list)
def create_ports_list(self, ports):
"""
Creates a list of ports from config data
Args:
item: A config data item representing the ports
"""
ports_list = ports["portsList"]
for port in ports_list:
port_value = port["port"]
self.ports_list.append(port_value)
# Set the number of ports
self.num_ports = len(self.ports_list)
def get_action_info(self, action_info):
"""
Extracts action_info
Args:
item: A config data item representing action info
"""
self.action_type = ACTION_TYPE[action_info["type"]]
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"])
logging.info("Training episodes have " + str(self.episode_steps) + " steps")
def reset_environment(self):
"""
# Resets environment using config data config data in order to build the environment configuration
"""
for item in self.config_data:
if item["itemType"] == "NODE":
# Reset a node's state (normal and reference)
self.reset_node(item)
elif item["itemType"] == "ACL_RULE":
# Create an ACL rule (these are cleared on reset, so just need to recreate them)
self.create_acl_rule(item)
else:
# Do nothing (bad formatting or not relevant to reset)
pass
# Reset the IER status so they are not running initially
# Green IERs
for ier_key, ier_value in self.green_iers.items():
ier_value.set_is_running(False)
# Red IERs
for ier_key, ier_value in self.red_iers.items():
ier_value.set_is_running(False)
def reset_node(self, item):
"""
Resets the statuses of a node
Args:
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_base_type = item["baseType"]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
node = self.nodes[node_id]
node_ref = self.nodes_reference[node_id]
# Reset the hardware state (common for all node types)
node.set_state(node_hardware_state)
node_ref.set_state(node_hardware_state)
if node_base_type == "ACTIVE":
# Active nodes have operating system state
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
elif node_base_type == "SERVICE":
# Service nodes have operating system state and list of services
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
# Update service states
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
service_state = SOFTWARE_STATE[service["state"]]
# Update node service state
node.set_service_state(service_protocol, service_state)
# Update reference node service state
node_ref.set_service_state(service_protocol, service_state)
else:
# Bad formatting
pass

View File

@@ -0,0 +1,224 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Implements reward function
"""
from common.enums import *
from nodes.active_node import ActiveNode
from nodes.service_node import ServiceNode
def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green_iers, red_iers, step_count, config_values):
"""
Compares the states of the initial and final nodes/links to get a reward
Args:
initial_nodes: The nodes before red and blue agents take effect
final_nodes: The nodes after red and blue agents take effect
reference_nodes: The nodes if there had been no red or blue effect
green_iers: The green IERs (should be running)
red_iers: Should be stopeed (ideally) by the blue agent
step_count: current step
"""
reward_value = 0
# For each node, compare operating state, o/s operating state, service states
for node_key, final_node in final_nodes.items():
initial_node = initial_nodes[node_key]
reference_node = reference_nodes[node_key]
# Operating State
reward_value += score_node_operating_state(final_node, initial_node, reference_node, config_values)
# Operating System State
if (isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode)):
reward_value += score_node_os_state(final_node, initial_node, reference_node, config_values)
# Service State
if (isinstance(final_node, ServiceNode)):
reward_value += score_node_service_state(final_node, initial_node, reference_node, config_values)
# Go through each red IER - penalise if it is running
for ier_key, ier_value in red_iers.items():
start_step = ier_value.get_start_step()
stop_step = ier_value.get_end_step()
if step_count >= start_step and step_count <= stop_step:
if ier_value.get_is_running():
reward_value += config_values.red_ier_running
# Go through each green IER - penalise if it's not running (weighted)
for ier_key, ier_value in green_iers.items():
start_step = ier_value.get_start_step()
stop_step = ier_value.get_end_step()
if step_count >= start_step and step_count <= stop_step:
if not ier_value.get_is_running():
reward_value += config_values.green_ier_blocked * ier_value.get_mission_criticality()
return reward_value
def score_node_operating_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating state of a node
Args:
final_node: The node after red and blue agents take effect
initial_node: The node before red and blue agents take effect
reference_node: The node if there had been no red or blue effect
"""
score = 0
final_node_operating_state = final_node.get_state()
initial_node_operating_state = initial_node.get_state()
reference_node_operating_state = reference_node.get_state()
if final_node_operating_state == reference_node_operating_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_node_operating_state == HARDWARE_STATE.ON:
if final_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_on
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_on
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.OFF:
if final_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_off
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_off
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
if final_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_resetting
elif final_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_resetting
else:
pass
else:
pass
return score
def score_node_os_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating system state of a node
Args:
final_node: The node after red and blue agents take effect
initial_node: The node before red and blue agents take effect
reference_node: The node if there had been no red or blue effect
"""
score = 0
final_node_os_state = final_node.get_os_state()
initial_node_os_state = initial_node.get_os_state()
reference_node_os_state = reference_node.get_os_state()
if final_node_os_state == reference_node_os_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_node_os_state == SOFTWARE_STATE.GOOD:
if final_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING:
if final_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED:
if final_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
else:
pass
else:
pass
return score
def score_node_service_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the service state(s) of a node
Args:
final_node: The node after red and blue agents take effect
initial_node: The node before red and blue agents take effect
reference_node: The node if there had been no red or blue effect
"""
score = 0
final_node_services = final_node.get_services()
initial_node_services = initial_node.get_services()
reference_node_services = reference_node.get_services()
for service_key, final_service in final_node_services.items():
reference_service = reference_node_services[service_key]
initial_service = initial_node_services[service_key]
if final_service.get_state() == reference_service.get_state():
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_service.get_state() == SOFTWARE_STATE.GOOD:
if final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_good
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed
else:
pass
else:
pass
return score

View File

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

132
PRIMAITE/links/link.py Normal file
View File

@@ -0,0 +1,132 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The link class
"""
from common.protocol import Protocol
from common.enums import *
class Link(object):
"""
Link class
"""
def __init__(self, _id, _bandwidth, _source_node_name, _dest_node_name, _services):
"""
Init
Args:
_id: The IER id
_bandwidth: The bandwidth of the link (bps)
_source_node_name: The name of the source node
_dest_node_name: The name of the destination node
_protocols: The protocols to add to the link
"""
self.id = _id
self.bandwidth = _bandwidth
self.source_node_name = _source_node_name
self.dest_node_name = _dest_node_name
self.protocol_list = []
# Add the default protocols
for protocol_name in _services:
self.add_protocol(protocol_name)
def add_protocol(self, _protocol):
"""
Adds a new protocol to the list of protocols on this link
Args:
_protocol: The protocol to be added (enum)
"""
self.protocol_list.append(Protocol(_protocol))
def get_id(self):
"""
Gets link ID
Returns:
Link ID
"""
return self.id
def get_source_node_name(self):
"""
Gets source node name
Returns:
Source node name
"""
return self.source_node_name
def get_dest_node_name(self):
"""
Gets destination node name
Returns:
Destination node name
"""
return self.dest_node_name
def get_bandwidth(self):
"""
Gets bandwidth of link
Returns:
Link bandwidth (bps)
"""
return self.bandwidth
def get_protocol_list(self):
"""
Gets list of protocols on this link
Returns:
List of protocols on this link
"""
return self.protocol_list
def get_current_load(self):
"""
Gets current total load on this link
Returns:
Total load on this link (bps)
"""
total_load = 0
for protocol in self.protocol_list:
total_load += protocol.get_load()
return total_load
def add_protocol_load(self, _protocol, _load):
"""
Adds a loading to a protocol on this link
Args:
_protocol: The protocol to load
_load: The amount to load (bps)
"""
for protocol in self.protocol_list:
if protocol.get_name() == _protocol:
protocol.add_load(_load)
else:
pass
def clear_traffic(self):
"""
Clears all traffic on this link
"""
for protocol in self.protocol_list:
protocol.clear_load()

View File

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

View File

@@ -0,0 +1,95 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
An Active Node (i.e. not an actuator)
"""
from nodes.node import Node
from common.enums import *
class ActiveNode(Node):
"""
Active Node class
"""
def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _config_values):
"""
Init
Args:
_id: The node ID
_name: The node name
_type: The node type (enum)
_priority: The node priority (enum)
_state: The node state (enum)
_ip_address: The node IP address
_os_state: The node Operating System state
"""
super().__init__(_id, _name, _type, _priority, _state, _config_values)
self.ip_address = _ip_address
self.os_state = _os_state
self.patching_count = 0
def set_ip_address(self, _ip_address):
"""
Sets IP address
Args:
_ip_address: IP address
"""
self.ip_address = _ip_address
def get_ip_address(self):
"""
Gets IP address
Returns:
IP address
"""
return self.ip_address
def set_os_state(self, _os_state):
"""
Sets operating system state
Args:
_os_state: Operating system state
"""
self.os_state = _os_state
if _os_state == SOFTWARE_STATE.PATCHING:
self.patching_count = self.config_values.os_patching_duration
def set_os_state_if_not_compromised(self, _os_state):
"""
Sets operating system state if the node is not compromised
Args:
_os_state: Operating system state
"""
if self.os_state != SOFTWARE_STATE.COMPROMISED:
self.os_state = _os_state
if _os_state == SOFTWARE_STATE.PATCHING:
self.patching_count = self.config_values.os_patching_duration
def get_os_state(self):
"""
Gets operating system state
Returns:
Operating system state
"""
return self.os_state
def update_os_patching_status(self):
"""
Updates operating system status based on patching cycle
"""
self.patching_count -= 1
if self.patching_count <= 0:
self.patching_count = 0
self.os_state = SOFTWARE_STATE.GOOD

176
PRIMAITE/nodes/node.py Normal file
View File

@@ -0,0 +1,176 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The base Node class
"""
from common.enums import *
class Node:
"""
Node class
"""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
"""
Init
Args:
_id: The node id
_name: The name of the node
_type: The type of the node
_priority: The priority of the node
_state: The state of the node
"""
self.id = _id
self.name = _name
self.type = _type
self.priority = _priority
self.operating_state = _state
self.resetting_count = 0
self.config_values = _config_values
def __repr__(self):
"""
Returns the name of the node
"""
return self.name
def set_id(self, _id):
"""
Sets the node ID
Args:
_id: The node ID
"""
self.id = _id
def get_id(self):
"""
Gets the node ID
Returns:
The node ID
"""
return self.id
def set_name(self, _name):
"""
Sets the node name
Args:
_name: The node name
"""
self.name = _name
def get_name(self):
"""
Gets the node name
Returns:
The node name
"""
return self.name
def set_type(self, _type):
"""
Sets the node type
Args:
_type: The node type
"""
self.type = _type
def get_type(self):
"""
Gets the node type
Returns:
The node type
"""
return self.type
def set_priority(self, _priority):
"""
Sets the node priority
Args:
_priority: The node priority
"""
self.priority = _priority
def get_priority(self):
"""
Gets the node priority
Returns:
The node priority
"""
return self.priority
def set_state(self, _state):
"""
Sets the node state
Args:
_state: The node state
"""
self.operating_state = _state
def get_state(self):
"""
Gets the node operating state
Returns:
The node operating state
"""
return self.operating_state
def turn_on(self):
"""
Sets the node state to ON
"""
self.operating_state = HARDWARE_STATE.ON
def turn_off(self):
"""
Sets the node state to OFF
"""
self.operating_state = HARDWARE_STATE.OFF
def reset(self):
"""
Sets the node state to Resetting and starts the reset count
"""
self.operating_state = HARDWARE_STATE.RESETTING
self.resetting_count = self.config_values.node_reset_duration
def update_resetting_status(self):
"""
Updates the resetting count
"""
self.resetting_count -= 1
if self.resetting_count <= 0:
self.resetting_count = 0
self.operating_state = HARDWARE_STATE.ON

View File

@@ -0,0 +1,104 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Defines node behaviour for PoL
"""
class NodeStateInstruction(object):
"""
The Node State Instruction class
"""
def __init__(self, _id, _start_step, _end_step, _node_id, _node_pol_type, _service_name, _state, _is_entry_node=False):
"""
Init
Args:
_id: The node state instruction id
_start_step: The start step of the instruction
_end_step: The end step of the instruction
_node_id: The id of the associated node
_node_pol_type: The pattern of life type
_service_name: The service name
_state: The state (node or service)
_is_entry_node: Indicator for entry node (default = False)
"""
self.id = _id
self.start_step = _start_step
self.end_step = _end_step
self.node_id = _node_id
self.node_pol_type = _node_pol_type
self.service_name = _service_name # Not used when not a service instruction
self.state = _state
self.is_entry_node = _is_entry_node
def get_start_step(self):
"""
Gets the start step
Returns:
The start step
"""
return self.start_step
def get_end_step(self):
"""
Gets the end step
Returns:
The end step
"""
return self.end_step
def get_node_id(self):
"""
Gets the node ID
Returns:
The node ID
"""
return self.node_id
def get_node_pol_type(self):
"""
Gets the node pattern of life type (enum)
Returns:
The node pattern of life type (enum)
"""
return self.node_pol_type
def get_service_name(self):
"""
Gets the service name
Returns:
The service name
"""
return self.service_name
def get_state(self):
"""
Gets the state (node or service)
Returns:
The state (node or service)
"""
return self.state
def get_is_entry_node(self):
"""
Informs of entry node
Returns:
True if entry node
"""
return self.is_entry_node

View File

@@ -0,0 +1,37 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The Passive Node class (i.e. an actuator)
"""
from nodes.node import Node
class PassiveNode(Node):
"""
The Passive Node class
"""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
"""
Init
Args:
_id: The node id
_name: The name of the node
_type: The type of the node
_priority: The priority of the node
_state: The state of the node
"""
# Pass through to Super for now
super().__init__(_id, _name, _type, _priority, _state, _config_values)
def get_ip_address(self):
"""
Gets the node IP address
Returns:
The node IP address
"""
# No concept of IP address for passive nodes for now
return ""

View File

@@ -0,0 +1,161 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
A Service Node (i.e. not an actuator)
"""
from nodes.active_node import ActiveNode
from common.enums import *
class ServiceNode(ActiveNode):
"""
ServiceNode class
"""
def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _config_values):
"""
Init
Args:
_id: The node id
_name: The name of the node
_type: The type of the node
_priority: The priority of the node
_state: The state of the node
_ipAddress: The IP address of the node
_osState: The operating system state of the node
"""
super().__init__(_id, _name, _type, _priority, _state, _ip_address, _os_state, _config_values)
self.services = {}
def add_service(self, _service):
"""
Adds a service to the node
Args:
_service: The service to add
"""
self.services[_service.get_name()] = _service
def get_services(self):
"""
Gets the dictionary of services on this node
Returns:
Dictionary of services on this node
"""
return self.services
def has_service(self, _protocol):
"""
Indicates whether a service is on a node
Returns:
True if service (protocol) is on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
return True
else:
pass
return False
def service_running(self, _protocol):
"""
Indicates whether a service is in a running state on the node
Returns:
True if service (protocol) is in a running state on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
if service_value.get_state() != SOFTWARE_STATE.PATCHING:
return True
else:
return False
else:
pass
return False
def service_is_overwhelmed(self, _protocol):
"""
Indicates whether a service is in an overwhelmed state on the node
Returns:
True if service (protocol) is in an overwhelmed state on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
if service_value.get_state() == SOFTWARE_STATE.OVERWHELMED:
return True
else:
return False
else:
pass
return False
def set_service_state(self, _protocol, _state):
"""
Sets the state of a service (protocol) on the node
Args:
_protocol: The service (protocol)
_state: The state value
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
# Can't set to compromised if you're in a patching state
if (_state == SOFTWARE_STATE.COMPROMISED and service_value.get_state() != SOFTWARE_STATE.PATCHING) or _state != SOFTWARE_STATE.COMPROMISED:
service_value.set_state(_state)
else:
# Do nothing
pass
if _state == SOFTWARE_STATE.PATCHING:
service_value.patching_count = self.config_values.service_patching_duration
else:
# Do nothing
pass
def set_service_state_if_not_compromised(self, _protocol, _state):
"""
Sets the state of a service (protocol) on the node if the operating state is not "compromised"
Args:
_protocol: The service (protocol)
_state: The state value
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
if service_value.get_state() != SOFTWARE_STATE.COMPROMISED:
service_value.set_state(_state)
if _state == SOFTWARE_STATE.PATCHING:
service_value.patching_count = self.config_values.service_patching_duration
def get_service_state(self, _protocol):
"""
Gets the state of a service
Returns:
The state of the service
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
return service_value.get_state()
def update_services_patching_status(self):
"""
Updates the patching counter for any service that are patching
"""
for service_key, service_value in self.services.items():
if service_value.get_state() == SOFTWARE_STATE.PATCHING:
service_value.reduce_patching_count()

1
PRIMAITE/pol/__init__.py Normal file
View File

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

226
PRIMAITE/pol/green_pol.py Normal file
View File

@@ -0,0 +1,226 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Implements Pattern of Life on the network (nodes and links)
"""
from networkx import shortest_path
from common.enums import *
from nodes.active_node import ActiveNode
from nodes.service_node import ServiceNode
_VERBOSE = False
def apply_iers(network, nodes, links, iers, acl, step):
"""
Applies IERs to the links (link pattern of life)
Args:
network: The network modelled in the environment
nodes: The nodes within the environment
links: The links within the environment
iers: The IERs to apply to the links
acl: The Access Control List
step: The step number
"""
if _VERBOSE:
print("Applying IERs")
# Go through each IER and check the conditions for it being applied
# If everything is in place, apply the IER protocol load to the relevant links
for ier_key, ier_value in iers.items():
start_step = ier_value.get_start_step()
stop_step = ier_value.get_end_step()
protocol = ier_value.get_protocol()
port = ier_value.get_port()
load = ier_value.get_load()
source_node_id = ier_value.get_source_node_id()
dest_node_id = ier_value.get_dest_node_id()
# Need to set the running status to false first for all IERs
ier_value.set_is_running(False)
source_valid = True
dest_valid = True
acl_block = False
if step >= start_step and step <= stop_step:
# continue --------------------------
# Get the source and destination node for this link
source_node = nodes[source_node_id]
dest_node = nodes[dest_node_id]
# 1. Check the source node situation
if source_node.get_type() == TYPE.SWITCH:
# It's a switch
if source_node.get_state() == HARDWARE_STATE.ON and source_node.get_os_state() != SOFTWARE_STATE.PATCHING:
source_valid = True
else:
# IER no longer valid
source_valid = False
elif source_node.get_type() == TYPE.ACTUATOR:
# It's an actuator
# TO DO
pass
else:
# It's not a switch or an actuator (so active node)
if source_node.get_state() == HARDWARE_STATE.ON and source_node.get_os_state() != SOFTWARE_STATE.PATCHING:
if source_node.has_service(protocol):
if source_node.service_running(protocol) and not source_node.service_is_overwhelmed(protocol):
source_valid = True
else:
source_valid = False
else:
# Do nothing - IER is not valid on this node
# (This shouldn't happen if the IER has been written correctly)
source_valid = False
else:
# Do nothing - IER no longer valid
source_valid = False
# 2. Check the dest node situation
if dest_node.get_type() == TYPE.SWITCH:
# It's a switch
if dest_node.get_state() == HARDWARE_STATE.ON and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING:
dest_valid = True
else:
# IER no longer valid
dest_valid = False
elif dest_node.get_type() == TYPE.ACTUATOR:
# It's an actuator
pass
else:
# It's not a switch or an actuator (so active node)
if dest_node.get_state() == HARDWARE_STATE.ON and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING:
if dest_node.has_service(protocol):
if dest_node.service_running(protocol) and not dest_node.service_is_overwhelmed(protocol):
dest_valid = True
else:
dest_valid = False
else:
# Do nothing - IER is not valid on this node
# (This shouldn't happen if the IER has been written correctly)
dest_valid = False
else:
# Do nothing - IER no longer valid
dest_valid = False
# 3. Check that the ACL doesn't block it
acl_block = acl.is_blocked(source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port)
if acl_block:
if _VERBOSE:
print("ACL block on source: " + source_node.get_ip_address() + ", dest: " + dest_node.get_ip_address() + ", protocol: " + protocol + ", port: " + port)
else:
if _VERBOSE:
print("No ACL block")
# Check whether both the source and destination are valid, and there's no ACL block
if source_valid and dest_valid and not acl_block:
# Load up the link(s) with the traffic
if _VERBOSE:
print("Source, Dest and ACL valid")
# Get the shortest path (i.e. nodes) between source and destination
path_node_list = shortest_path(network, source_node, dest_node)
path_node_list_length = len(path_node_list)
path_valid = True
# We might have a switch in the path, so check all nodes are operational
for node in path_node_list:
if node.get_state() != HARDWARE_STATE.ON or node.get_os_state() == SOFTWARE_STATE.PATCHING:
path_valid = False
if path_valid:
if _VERBOSE:
print("Applying IER to link(s)")
count = 0
link_capacity_exceeded = False
# Check that the link capacity is not exceeded by the new load
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
link_id = edge_dict[0].get('id')
link = links[link_id]
# Check whether the new load exceeds the bandwidth
if (link.get_current_load() + load) > link.get_bandwidth():
link_capacity_exceeded = True
if _VERBOSE:
print("Link capacity exceeded")
pass
count+=1
# Check whether the link capacity for any links on this path have been exceeded
if link_capacity_exceeded == False:
# Now apply the new loads to the links
count = 0
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
link_id = edge_dict[0].get('id')
link = links[link_id]
# Add the load from this IER
link.add_protocol_load(protocol, load)
count+=1
# This IER is now valid, so set it to running
ier_value.set_is_running(True)
else:
# One of the nodes is not operational
if _VERBOSE:
print("Path not valid - one or more nodes not operational")
pass
else:
if _VERBOSE:
print("Source, Dest or ACL were not valid")
pass
# ------------------------------------
else:
# Do nothing - IER no longer valid
pass
def apply_node_pol(nodes, node_pol, step):
"""
Applies node pattern of life
Args:
nodes: The nodes within the environment
node_pol: The node pattern of life to apply
step: The step number
"""
if _VERBOSE:
print("Applying Node PoL")
for key, node_instruction in node_pol.items():
start_step = node_instruction.get_start_step()
stop_step = node_instruction.get_end_step()
node_id = node_instruction.get_node_id()
node_pol_type = node_instruction.get_node_pol_type()
service_name = node_instruction.get_service_name()
state = node_instruction.get_state()
if step >= start_step and step <= stop_step:
# continue --------------------------
node = nodes[node_id]
if node_pol_type == NODE_POL_TYPE.OPERATING:
# Change operating state
node.set_state(state)
elif node_pol_type == NODE_POL_TYPE.OS:
# Change OS state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.set_os_state_if_not_compromised(state)
else:
# Change a service state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ServiceNode):
node.set_service_state_if_not_compromised(service_name, state)
else:
# PoL is not valid in this time step
pass

147
PRIMAITE/pol/ier.py Normal file
View File

@@ -0,0 +1,147 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Information Exchange Requirements for APE
Used to represent an information flow from source to destination
"""
class IER(object):
"""
Information Exchange Requirement class
"""
def __init__(self, _id, _start_step, _end_step, _load, _protocol, _port, _source_node_id, _dest_node_id, _mission_criticality, _running=False):
"""
Init
Args:
_id: The IER id
_start_step: The step when this IER should start
_end_step: The step when this IER should end
_load: The load this IER should put on a link (bps)
_protocol: The protocol of this IER
_port: The port this IER runs on
_source_node_id: The source node ID
_dest_node_id: The destination node ID
_mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
_running: Indicates whether the IER is currently running
"""
self.id = _id
self.start_step = _start_step
self.end_step = _end_step
self.source_node_id = _source_node_id
self.dest_node_id = _dest_node_id
self.load = _load
self.protocol = _protocol
self.port = _port
self.mission_criticality = _mission_criticality
self.running = _running
def get_id(self):
"""
Gets IER ID
Returns:
IER ID
"""
return self.id
def get_start_step(self):
"""
Gets IER start step
Returns:
IER start step
"""
return self.start_step
def get_end_step(self):
"""
Gets IER end step
Returns:
IER end step
"""
return self.end_step
def get_load(self):
"""
Gets IER load
Returns:
IER load
"""
return self.load
def get_protocol(self):
"""
Gets IER protocol
Returns:
IER protocol
"""
return self.protocol
def get_port(self):
"""
Gets IER port
Returns:
IER port
"""
return self.port
def get_source_node_id(self):
"""
Gets IER source node ID
Returns:
IER source node ID
"""
return self.source_node_id
def get_dest_node_id(self):
"""
Gets IER destination node ID
Returns:
IER destination node ID
"""
return self.dest_node_id
def get_is_running(self):
"""
Informs whether the IER is currently running
Returns:
True if running
"""
return self.running
def set_is_running(self, _value):
"""
Sets the running state of the IER
Args:
_value: running status
"""
self.running = _value
def get_mission_criticality(self):
"""
Gets the IER mission criticality (used in the reward function)
Returns:
Mission criticality value (0 lowest to 5 highest)
"""
return self.mission_criticality

View File

@@ -0,0 +1,272 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Implements Pattern of Life on the network (nodes and links) resulting from the red agent attack
"""
from networkx import shortest_path
from common.enums import *
from nodes.active_node import ActiveNode
from nodes.service_node import ServiceNode
_VERBOSE = False
def apply_red_agent_iers(network, nodes, links, iers, acl, step):
"""
Applies IERs to the links (link pattern of life) resulting from red agent attack
Args:
network: The network modelled in the environment
nodes: The nodes within the environment
links: The links within the environment
iers: The red agent IERs to apply to the links
acl: The Access Control List
step: The step number
"""
# Go through each IER and check the conditions for it being applied
# If everything is in place, apply the IER protocol load to the relevant links
for ier_key, ier_value in iers.items():
start_step = ier_value.get_start_step()
stop_step = ier_value.get_end_step()
protocol = ier_value.get_protocol()
port = ier_value.get_port()
load = ier_value.get_load()
source_node_id = ier_value.get_source_node_id()
dest_node_id = ier_value.get_dest_node_id()
# Need to set the running status to false first for all IERs
ier_value.set_is_running(False)
source_valid = True
dest_valid = True
acl_block = False
if step >= start_step and step <= stop_step:
# continue --------------------------
# Get the source and destination node for this link
source_node = nodes[source_node_id]
dest_node = nodes[dest_node_id]
# 1. Check the source node situation
if source_node.get_type() == TYPE.SWITCH:
# It's a switch
if source_node.get_state() == HARDWARE_STATE.ON:
source_valid = True
else:
# IER no longer valid
source_valid = False
elif source_node.get_type() == TYPE.ACTUATOR:
# It's an actuator
# TO DO
pass
else:
# It's not a switch or an actuator (so active node)
if source_node.get_state() == HARDWARE_STATE.ON:
if source_node.has_service(protocol):
# Red agents IERs can only be valid if the source service is in a compromised state
if source_node.get_service_state(protocol) == SOFTWARE_STATE.COMPROMISED:
source_valid = True
else:
source_valid = False
else:
# Do nothing - IER is not valid on this node
# (This shouldn't happen if the IER has been written correctly)
source_valid = False
else:
# Do nothing - IER no longer valid
source_valid = False
# 2. Check the dest node situation
if dest_node.get_type() == TYPE.SWITCH:
# It's a switch
if dest_node.get_state() == HARDWARE_STATE.ON:
dest_valid = True
else:
# IER no longer valid
dest_valid = False
elif dest_node.get_type() == TYPE.ACTUATOR:
# It's an actuator
pass
else:
# It's not a switch or an actuator (so active node)
if dest_node.get_state() == HARDWARE_STATE.ON:
if dest_node.has_service(protocol):
# We don't care what state the destination service is in for an IER
dest_valid = True
else:
# Do nothing - IER is not valid on this node
# (This shouldn't happen if the IER has been written correctly)
dest_valid = False
else:
# Do nothing - IER no longer valid
dest_valid = False
# 3. Check that the ACL doesn't block it
acl_block = acl.is_blocked(source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port)
if acl_block:
if _VERBOSE:
print("ACL block on source: " + source_node.get_ip_address() + ", dest: " + dest_node.get_ip_address() + ", protocol: " + protocol + ", port: " + port)
else:
if _VERBOSE:
print("No ACL block")
# Check whether both the source and destination are valid, and there's no ACL block
if source_valid and dest_valid and not acl_block:
# Load up the link(s) with the traffic
if _VERBOSE:
print("Source, Dest and ACL valid")
# Get the shortest path (i.e. nodes) between source and destination
path_node_list = shortest_path(network, source_node, dest_node)
path_node_list_length = len(path_node_list)
path_valid = True
# We might have a switch in the path, so check all nodes are operational
# We're assuming here that red agents can get past switches that are patching
for node in path_node_list:
if node.get_state() != HARDWARE_STATE.ON:
path_valid = False
if path_valid:
if _VERBOSE:
print("Applying IER to link(s)")
count = 0
link_capacity_exceeded = False
# Check that the link capacity is not exceeded by the new load
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
link_id = edge_dict[0].get('id')
link = links[link_id]
# Check whether the new load exceeds the bandwidth
if (link.get_current_load() + load) > link.get_bandwidth():
link_capacity_exceeded = True
if _VERBOSE:
print("Link capacity exceeded")
pass
count+=1
# Check whether the link capacity for any links on this path have been exceeded
if link_capacity_exceeded == False:
# Now apply the new loads to the links
count = 0
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
link_id = edge_dict[0].get('id')
link = links[link_id]
# Add the load from this IER
link.add_protocol_load(protocol, load)
count+=1
# This IER is now valid, so set it to running
ier_value.set_is_running(True)
if _VERBOSE:
print("Red IER was allowed to run in step " + str(step))
else:
# One of the nodes is not operational
if _VERBOSE:
print("Path not valid - one or more nodes not operational")
pass
else:
if _VERBOSE:
print("Red IER was NOT allowed to run in step " + str(step))
print("Source, Dest or ACL were not valid")
pass
# ------------------------------------
else:
# Do nothing - IER no longer valid
pass
pass
def apply_red_agent_node_pol(nodes, iers, node_pol, step):
"""
Applies node pattern of life
Args:
nodes: The nodes within the environment
iers: The red agent IERs
node_pol: The red agent node pattern of life to apply
step: The step number
"""
if _VERBOSE:
print("Applying Node Red Agent PoL")
for key, node_instruction in node_pol.items():
start_step = node_instruction.get_start_step()
stop_step = node_instruction.get_end_step()
node_id = node_instruction.get_node_id()
node_pol_type = node_instruction.get_node_pol_type()
service_name = node_instruction.get_service_name()
state = node_instruction.get_state()
is_entry_node = node_instruction.get_is_entry_node()
if step >= start_step and step <= stop_step:
# continue --------------------------
node = nodes[node_id]
# for the red agent, either:
# 1. the node has to be an entry node, or
# 2. there is a red IER relevant to that service entering the node with a running status of True
red_ier_incoming = is_red_ier_incoming(node, iers, node_pol_type)
if is_entry_node or red_ier_incoming:
if node_pol_type == NODE_POL_TYPE.OPERATING:
# Change operating state
node.set_state(state)
elif node_pol_type == NODE_POL_TYPE.OS:
# Change OS state
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.set_os_state(state)
else:
# Change a service state
if isinstance(node, ServiceNode):
node.set_service_state(service_name, state)
else:
if _VERBOSE:
print("Node Red Agent PoL not allowed - not entry node, or running IER not present")
else:
# PoL is not valid in this time step
pass
def is_red_ier_incoming(node, iers, node_pol_type):
node_id = node.get_id()
for ier_key, ier_value in iers.items():
if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id:
if node_pol_type == NODE_POL_TYPE.OPERATING or node_pol_type == NODE_POL_TYPE.OS:
# It's looking to change operating state or O/S state, so valid
return True
elif node_pol_type == NODE_POL_TYPE.SERVICE:
# Check if the service is present on the node and running
ier_protocol = ier_value.get_protocol()
if isinstance(node, ServiceNode):
if node.has_service(ier_protocol):
if node.service_running(ier_protocol):
# Matching service is present and running, so valid
return True
else:
# Service is present, but not running
return False
else:
# Service is not present
return False
else:
# Not a service node
return False
else:
# Shouldn't get here - instruction type is undefined
return False
else:
# The IER destination is not this node, or the IER is not running
return False

105
PRIMAITE/tests/test_acl.py Normal file
View File

@@ -0,0 +1,105 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Used to tes the ACL functions
"""
from acl.acl_rule import ACLRule
from acl.access_control_list import AccessControlList
def test_acl_address_match_1():
"""
Test that matching IP addresses produce True
"""
acl = AccessControlList()
rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
def test_acl_address_match_2():
"""
Test that mismatching IP addresses produce False
"""
acl = AccessControlList()
rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False
def test_acl_address_match_3():
"""
Test the ANY condition for source IP addresses produce True
"""
acl = AccessControlList()
rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
def test_acl_address_match_4():
"""
Test the ANY condition for dest IP addresses produce True
"""
acl = AccessControlList()
rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
def test_check_acl_block_affirmative():
"""
Test the block function (affirmative)
"""
# Create the Access Control List
acl = AccessControlList()
# Create a rule
acl_rule_permission = "ALLOW"
acl_rule_source = "192.168.1.1"
acl_rule_destination = "192.168.1.2"
acl_rule_protocol = "TCP"
acl_rule_port = "80"
acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False
def test_check_acl_block_negative():
"""
Test the block function (negative)
"""
# Create the Access Control List
acl = AccessControlList()
# Create a rule
acl_rule_permission = "DENY"
acl_rule_source = "192.168.1.1"
acl_rule_destination = "192.168.1.2"
acl_rule_protocol = "TCP"
acl_rule_port = "80"
acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == True
def test_rule_hash():
"""
Test the rule hash
"""
# Create the Access Control List
acl = AccessControlList()
rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
hash_value_local = hash(rule)
hash_value_remote = acl.get_dictionary_hash("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
assert hash_value_local == hash_value_remote

View File

View File

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

View File

@@ -0,0 +1,69 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The Transaction class
"""
class Transaction(object):
"""
Transaction class
"""
def __init__(self, _timestamp, _agent_identifier, _episode_number, _step_number):
"""
Init
Args:
_timestamp: The time this object was created
_agent_identifier: An identifier for the agent in use
_episode_number: The episode number
_step_number: The step number
"""
self.timestamp = _timestamp
self.agent_identifier = _agent_identifier
self.episode_number = _episode_number
self.step_number = _step_number
def set_obs_space_pre(self, _obs_space_pre):
"""
Sets the observation space (pre)
Args:
_obs_space_pre: The observation space before any actions are taken
"""
self.obs_space_pre = _obs_space_pre
def set_obs_space_post(self, _obs_space_post):
"""
Sets the observation space (post)
Args:
_obs_space_post: The observation space after any actions are taken
"""
self.obs_space_post = _obs_space_post
def set_reward(self, _reward):
"""
Sets the reward
Args:
_reward: The reward value
"""
self.reward = _reward
def set_action_space(self, _action_space):
"""
Sets the action space
Args:
_action_space: The action space invoked by the agent
"""
self.action_space = _action_space

View File

@@ -0,0 +1,104 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Writes the Transaction log list out to file for evaluation to utilse
"""
import csv
import logging
import os.path
from datetime import datetime
from transactions.transaction import Transaction
def turn_action_space_to_array(_action_space):
"""
Turns action space into a string array so it can be saved to csv
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
def turn_obs_space_to_array(_obs_space, _obs_assets, _obs_features):
"""
Turns observation space into a string array so it can be saved to csv
Args:
_obs_space: The observation space
_obs_assets: The number of assets (i.e. nodes or links) in the observation space
_obs_features: The number of features associated with the asset
"""
return_array = []
for x in range(_obs_assets):
for y in range(_obs_features):
return_array.append(str(_obs_space[x][y]))
return return_array
def write_transaction_to_file(_transaction_list):
"""
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
"""
# 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]
action_length = template_transation.action_space.size
obs_assets = template_transation.obs_space_post.shape[0]
obs_features = template_transation.obs_space_post.shape[1]
# Create the action space headers array
action_header = []
for x in range(action_length):
action_header.append('AS_' + str(x))
# Create the observation space headers array
obs_header_initial = []
obs_header_new = []
for x in range(obs_assets):
for y in range(obs_features):
obs_header_initial.append('OSI_' + str(x) + '_' + str(y))
obs_header_new.append('OSN_' + str(x) + '_' + str(y))
# 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"
csv_file = open(filename, 'w', encoding='UTF8', newline='')
csv_writer = csv.writer(csv_file)
csv_writer.writerow(header)
for transaction in _transaction_list:
csv_data = [str(transaction.timestamp), str(transaction.episode_number), str(transaction.step_number), str(transaction.reward)]
csv_data = csv_data + turn_action_space_to_array(transaction.action_space) + \
turn_obs_space_to_array(transaction.obs_space_pre, obs_assets, obs_features) + \
turn_obs_space_to_array(transaction.obs_space_post, obs_assets, obs_features)
csv_writer.writerow(csv_data)
csv_file.close()
except Exception as e:
logging.error("Could not save the transaction file")
logging.error("Exception occured", exc_info=True)

View File

@@ -1,20 +1 @@
# Introduction
TODO: Give a short introduction of your project. Let this section explain the objectives or the motivation behind this project.
# Getting Started
TODO: Guide users through getting your code up and running on their own system. In this section you can talk about:
1. Installation process
2. Software dependencies
3. Latest releases
4. API references
# Build and Test
TODO: Describe and show how to build your code and run the tests.
# Contribute
TODO: Explain how other users and developers can contribute to make your code better.
If you want to learn more about creating good readme files then refer the following [guidelines](https://docs.microsoft.com/en-us/azure/devops/repos/git/create-a-readme?view=azure-devops). You can also seek inspiration from the below readme files:
- [ASP.NET Core](https://github.com/aspnet/Home)
- [Visual Studio Code](https://github.com/Microsoft/vscode)
- [Chakra Core](https://github.com/Microsoft/ChakraCore)
# PrimAITE

27
setup.py Normal file
View File

@@ -0,0 +1,27 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Setup
"""
from setuptools import find_packages, setup
setup(
name="primaite",
maintainer="QinetiQ Training and Simulation Ltd",
url="https://github.com/qtsl/PrimAITE",
description="A primary-level simulation tool",
python_requires=">=3.7",
version="1.0.0",
install_requires=[
"gym==0.21.0",
"matplotlib == 3.5.2",
"networkx == 2.6.3",
"numpy == 1.21.1",
"stable_baselines3 == 1.6.0",
"pandas == 1.1.5",
"pyyaml == 6.0",
"typing-extensions == 4.2.0",
"torch == 1.12.0"
],
packages=find_packages()
)