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

Merge remote-tracking branch 'origin/dev' into release/3.0.0b7

Browse Source
This commit is contained in:
Marek Wolan
2024-03-15 16:13:50 +00:00
parent 317fbdbb9c d5fb2f2309
commit a8c98c28f1
192 changed files with 10846 additions and 3590 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
CHANGELOG.md
View File

@@ -6,27 +6,35 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
- Made requests fail to reach their target if the node is off
- Added responses to requests
- Made environment reset completely recreate the game object.
- Changed the red agent in the data manipulation scenario to randomly choose client 1 or client 2 to start its attack.
- Changed the data manipulation scenario to include a second green agent on client 1.
- Refactored actions and observations to be configurable via object name, instead of UUID.
- Fixed a bug where ACL rules were not resetting on episode reset.
- Fixed a bug where blue agent's ACL actions were being applied against the wrong IP addresses
- Fixed a bug where deleted files and folders did not reset correctly on episode reset.
- Fixed a bug where service health status was using the actual health state instead of the visible health state
- Fixed a bug where the database file health status was using the incorrect value for negative rewards
- Fixed a bug preventing file actions from reaching their intended file
- Made database patch correctly take 2 timesteps instead of being immediate
- Made database patch only possible when the software is compromised or good, it's no longer possible when the software is OFF or RESETTING
- Temporarily disable the blue agent file delete action due to crashes. This issue is resolved in another branch that will be merged into dev soon.
- Fix a bug where ACLs were not showing up correctly in the observation space.
- Added a notebook which explains Data manipulation scenario, demonstrates the attack, and shows off blue agent's action space, observation space, and reward function.
- Made packet capture and system logging optional (off by default). To turn on, change the io_settings.save_pcap_logs and io_settings.save_sys_logs settings in the config.
- Made observation space flattening optional (on by default). To turn off for an agent, change the agent_settings.flatten_obs setting in the config.
- Fixed an issue where the data manipulation attack was triggered at episode start.
- Fixed a bug where FTP STOR stored an additional copy on the client machine's filesystem
- Fixed a bug where the red agent acted to early
- Fixed the order of service health state
- Fixed an issue where starting a node didn't start the services on it
- Made observation space flattening optional (on by default). To turn off for an agent, change the `agent_settings.flatten_obs` setting in the config.
- Added support for SQL INSERT command.
- Added ability to log each agent's action choices in each step to a JSON file.
### Bug Fixes
- ACL rules were not resetting on episode reset.
- ACLs were not showing up correctly in the observation space.
- Blue agent's ACL actions were being applied against the wrong IP addresses
- Deleted files and folders did not reset correctly on episode reset.
- Service health status was using the actual health state instead of the visible health state
- Database file health status was using the incorrect value for negative rewards
- Preventing file actions from reaching their intended file
- The data manipulation attack was triggered at episode start.
- FTP STOR stored an additional copy on the client machine's filesystem
- The red agent acted to early
- Order of service health state
- Starting a node didn't start the services on it
- Fixed an issue where the services were still able to run even though the node the service is installed on is turned off
@@ -46,8 +54,12 @@ a Service/Application another machine.
SessionManager.
- Permission System - each action can define criteria that will be used to permit or deny agent actions.
- File System - ability to emulate a node's file system during a simulation
- Example notebooks - There is currently 1 jupyter notebook which walks through using PrimAITE
1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP)
- Example notebooks - There are 5 jupyter notebook which walk through using PrimAITE
1. Training a Stable Baselines 3 agent
2. Training a single agent system using Ray RLLib
3. Training a multi-agent system Ray RLLib
4. Data manipulation end to end demonstration
5. Data manipulation scenario with customised red agents
- Database:
- `DatabaseClient` and `DatabaseService` created to allow emulation of database actions
- Ability for `DatabaseService` to backup its data to another server via FTP and restore data from backup
@@ -57,7 +69,6 @@ SessionManager.
- DNS Services: `DNSClient` and `DNSServer`
- FTP Services: `FTPClient` and `FTPServer`
- HTTP Services: `WebBrowser` to simulate a web client and `WebServer`
- Fixed an issue where the services were still able to run even though the node the service is installed on is turned off
- NTP Services: `NTPClient` and `NTPServer`
- **RouterNIC Class**: Introduced a new class `RouterNIC`, extending the standard `NIC` functionality. This class is specifically designed for router operations, optimizing the processing and routing of network traffic.
- **Custom Layer-3 Processing**: The `RouterNIC` class includes custom handling for network frames, bypassing standard Node NIC's Layer 3 broadcast/unicast checks. This allows for more efficient routing behavior in network scenarios where router-specific frame processing is required.
@@ -81,7 +92,21 @@ SessionManager.
- `AirSpace` class to simulate wireless communications, managing wireless interfaces and facilitating the transmission of frames within specified frequencies.
- `AirSpaceFrequency` enum for defining standard wireless frequencies, including 2.4 GHz and 5 GHz bands, to support realistic wireless network simulations.
- `WirelessRouter` class, extending the `Router` class, to incorporate wireless networking capabilities alongside traditional wired connections. This class allows the configuration of wireless access points with specific IP settings and operating frequencies.
- Documentation Updates:
- Examples include how to set up PrimAITE session via config
- Examples include how to create nodes and install software via config
- Examples include how to set up PrimAITE session via Python
- Examples include how to create nodes and install software via Python
- Added missing ``DoSBot`` documentation page
- Added diagrams where needed to make understanding some things easier
- Templated parts of the documentation to prevent unnecessary repetition and for easier maintaining of documentation
- Separated documentation pages of some items i.e. client and server software were on the same pages - which may make things confusing
- Configuration section at the bottom of the software pages specifying the configuration options available (and which ones are optional)
- Ability to add ``Firewall`` node via config
- Ability to add ``Router`` routes via config
- Ability to add ``Router``/``Firewall`` ``ACLRule`` via config
- NMNE capturing capabilities to `NetworkInterface` class for detecting and logging Malicious Network Events.
- New `nmne_config` settings in the simulation configuration to enable NMNE capturing and specify keywords such as "DELETE".
### Changed
- Integrated the RouteTable into the Routers frame processing.
@@ -93,7 +118,9 @@ SessionManager.
- Refactored all tests to utilise new `Node` subclasses (`Computer`, `Server`, `Router`, `Switch`) instead of creating generic `Node` instances and manually adding network interfaces. This change aligns test setups more closely with the intended use cases and hierarchies within the network simulation framework.
- Updated all tests to employ the `Network()` class for managing nodes and their connections, ensuring a consistent and structured approach to setting up network topologies in testing scenarios.
- **ACLRule Wildcard Masking**: Updated the `ACLRule` class to support IP ranges using wildcard masking. This enhancement allows for more flexible and granular control over traffic filtering, enabling the specification of broader or more specific IP address ranges in ACL rules.
- Updated `NetworkInterface` documentation to reflect the new NMNE capturing features and how to use them.
- Integration of NMNE capturing functionality within the `NicObservation` class.
- Changed blue action set to enable applying node scan, reset, start, and shutdown to every host in data manipulation scenario
### Removed
- Removed legacy simulation modules: `acl`, `common`, `environment`, `links`, `nodes`, `pol`
@@ -103,7 +130,7 @@ SessionManager.
### Fixed
- Addressed network transmission issues that previously allowed ARP requests to be incorrectly routed and repeated across different subnets. This fix ensures ARP requests are correctly managed and confined to their appropriate network segments.
- Resolved problems in `Node` and its subclasses where the default gateway configuration was not properly utilized for communications across different subnets. This correction ensures that nodes effectively use their configured default gateways for outbound communications to other network segments, thereby enhancing the network's routing functionality and reliability.
- Network Interface Port name/num being set properly for sys log and PCAP output.
## [2.0.0] - 2023-07-26

2
docs/Makefile
View File

@@ -6,7 +6,7 @@ SPHINXBUILD ?= sphinx-build
SOURCEDIR = .
BUILDDIR = _build
AUTOSUMMARY="source\_autosummary"
AUTOSUMMARY="source/_autosummary"
# Remove command is different depending on OS
ifdef OS

BIN
docs/_static/firewall_acl.png vendored Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 35 KiB

BIN
docs/_static/notebooks/extensions.png vendored Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 68 KiB

BIN
docs/_static/notebooks/install_extensions.png vendored Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 193 KiB

BIN
docs/_static/switched_p2p_network.png vendored Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 9.0 KiB

2
docs/api.rst
View File

@@ -1,3 +1,5 @@
:orphan:
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK

35
docs/conf.py
View File

@@ -10,12 +10,12 @@ import datetime
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
import os
import sys
from typing import Any
import furo # noqa
sys.path.insert(0, os.path.abspath("../"))
# -- Project information -----------------------------------------------------
year = datetime.datetime.now().year
project = "PrimAITE"
@@ -28,6 +28,11 @@ with open("../src/primaite/VERSION", "r") as file:
# The full version, including alpha/beta/rc tags
release = version
# set global variables
rst_prolog = f"""
.. |VERSION| replace:: {release}
"""
html_title = f"{project} v{release} docs"
# -- General configuration ---------------------------------------------------
@@ -45,13 +50,35 @@ extensions = [
"sphinx_copybutton", # Adds a copy button to code blocks
]
templates_path = ["_templates"]
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
exclude_patterns = [
"_build",
"Thumbs.db",
".DS_Store",
]
# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
html_theme = "furo"
html_static_path = ["_static"]
html_theme_options = {"globaltoc_collapse": True, "globaltoc_maxdepth": 2}
html_copy_source = False
def replace_token(app: Any, docname: Any, source: Any):
"""Replaces a token from the list of tokens."""
result = source[0]
for key in app.config.tokens:
result = result.replace(key, app.config.tokens[key])
source[0] = result
tokens = {"{VERSION}": release} # Token VERSION is replaced by the value of the PrimAITE version in the version file
"""Dict containing the tokens that need to be replaced in documentation."""
def setup(app: Any):
"""Custom setup for sphinx."""
app.add_config_value("tokens", {}, True)
app.connect("source-read", replace_token)

3
docs/index.rst
View File

@@ -105,9 +105,12 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE!
source/getting_started
source/primaite_session
source/example_notebooks
source/simulation
source/game_layer
source/config
source/environment
source/customising_scenarios
.. toctree::
:caption: Developer information:

128
docs/source/config.rst
View File

@@ -1,102 +1,40 @@
Primaite v3 config
******************
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
PrimAITE |VERSION| Configuration
********************************
PrimAITE uses a single configuration file to define everything needed to train and evaluate an RL policy in a custom cybersecurity scenario. This includes the configuration of the network, the scripted or trained agents that interact with the network, as well as settings that define how to perform training in Stable Baselines 3 or Ray RLLib.
The entire config is used by the ``PrimaiteSession`` object for users who wish to let PrimAITE handle the agent definition and training. If you wish to define custom agents and control the training loop yourself, you can use the config with the ``PrimaiteGame``, and ``PrimaiteGymEnv`` objects instead. That way, only the network configuration and agent setup parts of the config are used, and the training section is ignored.
Example Configuration Hierarchy
###############################
The top level configuration items in a configuration file is as follows
.. code-block:: yaml
training_config:
...
io_settings:
...
game:
...
agents:
...
simulation:
...
These are expanded upon in the Configurable items section below
Configurable items
==================
##################
``training_config``
-------------------
This section allows selecting which training framework and algorithm to use, and set some training hyperparameters.
.. toctree::
:maxdepth: 1
``io_settings``
---------------
This section configures how PrimAITE saves data during simulation and training.
**save_final_model**: Only used if training with PrimaiteSession, if true, the policy will be saved after the final training iteration.
**save_checkpoints**: Only used if training with PrimaiteSession, if true, the policy will be saved periodically during training.
**checkpoint_interval**: Only used if training with PrimaiteSession and if ``save_checkpoints`` is true. Defines how often to save the policy during training.
**save_logs**: *currently unused*.
**save_transactions**: *currently unused*.
**save_tensorboard_logs**: *currently unused*.
**save_step_metadata**: Whether to save the RL agents' action, environment state, and other data at every single step.
**save_pcap_logs**: Whether to save pcap files of all network traffic during the simulation.
**save_sys_logs**: Whether to save system logs from all nodes during the simulation.
``game``
--------
This section defines high-level settings that apply across the game, currently it's used to help shape the action and observation spaces by restricting which ports and internet protocols should be considered. Here, users can also set the maximum number of steps in an episode.
``agents``
----------
Agents can be scripted (deterministic and stochastic), or controlled by a reinforcement learning algorithm. Not to be confused with an RL agent, the term agent here is used to refer to an entity that sends requests to the simulated network. In this part of the config, each agent's action space, observation space, and reward function can be defined. All three are defined in a modular way.
**type**: Specifies which class should be used for the agent. ``ProxyAgent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``RedDatabaseCorruptingAgent`` and ``GreenWebBrowsingAgent`` generate their own behaviour.
**team**: Specifies if the agent is malicious (RED), benign (GREEN), or defensive (BLUE). Currently this value is not used for anything.
**observation space:**
* ``type``: selects which python class from the ``primaite.game.agent.observation`` module is used for the overall observation structure.
* ``options``: allows configuring the chosen observation type. The ``UC2BlueObservation`` should be used for RL Agents.
* ``num_services_per_node``, ``num_folders_per_node``, ``num_files_per_folder``, ``num_nics_per_node`` all define the shape of the observation space. The size and shape of the obs space must remain constant, but the number of files, folders, ACL rules, and other components can change within an episode. Therefore padding is performed and these options set the size of the obs space.
* ``nodes``: list of nodes that will be present in this agent's observation space. The ``node_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config. Each node can also be configured with services, and files that should be monitored.
* ``links``: list of links that will be present in this agent's observation space. The ``link_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config.
* ``acl``: configure how the agent reads the access control list on the router in the simulation. ``router_node_ref`` is for selecting which router's ACL table should be used. ``ip_address_order`` sets the encoding of ip addresses as integers within the observation space.
**action space:**
The action space is configured to be made up of individual action types. Once configured, the agent can select an action type and some optional action parameters at every step. For example: The ``NODE_SERVICE_SCAN`` action takes the parameters ``node_id`` and ``service_id``.
Description of configurable items:
* ``action_list``: a list of action modules. The options are listed in the ``primaite.game.agent.actions`` module.
* ``action_map``: (optional). Restricts the possible combinations of action type / action parameter values to reduce the overall size of the action space. By default, every possible combination of actions and parameters will be assigned an integer for the agent's ``MultiDiscrete`` action space. Instead, the ``action_map`` allows you to list the actions corresponding to each integer in the ``MultiDiscrete`` space.
* ``options``: Options that apply too all action components.
* ``nodes``: list the nodes that the agent can act on, the order of this list defines the mapping between nodes and ``node_id`` integers.
* ``max_folders_per_node``, ``max_files_per_folder``, ``max_services_per_node``, ``max_nics_per_node``, ``max_acl_rules`` all are used to define the size of the action space.
**reward function:**
Similar to action space, this is defined as a list of components.
Description of configurable items:
* ``reward_components`` a list of reward components from the ``primaite.game.agent.reward`` module.
* ``weight``: relative importance of this reward component. The total reward for a step is a weighted sum of all reward components.
* ``options``: list of options passed to the reward component during initialisation, the exact options required depend on the reward component.
**agent_settings**:
Settings passed to the agent during initialisation. These depend on the agent class.
Reinforcement learning agents use the ``ProxyAgent`` class, they accept these agent settings:
**flatten_obs**: If true, gymnasium flattening will be performed on the observation space before sending to the agent. Set this to true if your agent does not support nested observation spaces.
``simulation``
--------------
In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents.
At the top level of the network are ``nodes`` and ``links``.
**nodes:**
* ``type``: one of ``router``, ``switch``, ``computer``, or ``server``, this affects what other sub-options should be defined.
* ``hostname`` - a non-unique name used for logging and outputs.
* ``num_ports`` (optional, routers and switches only): number of network interfaces present on the device.
* ``ports`` (optional, routers and switches only): configuration for each network interface, including IP address and subnet mask.
* ``acl`` (Router only): Define the ACL rules at each index of the ACL on the router. the possible options are: ``action`` (PERMIT or DENY), ``src_port``, ``dst_port``, ``protocol``, ``src_ip``, ``dst_ip``. Any options left blank default to none which usually means that it will apply across all options. For example leaving ``src_ip`` blank will apply the rule to all IP addresses.
* ``services`` (computers and servers only): a list of services to install on the node. They must define a ``ref``, ``type``, and ``options`` that depend on which ``type`` was selected.
* ``applications`` (computer and servers only): Similar to services. A list of application to install on the node.
* ``network_interfaces`` (computers and servers only): If the node has multiple networking devices, the second, third, fourth, etc... must be defined here with an ``ip_address`` and ``subnet_mask``.
**links:**
* ``ref``: unique identifier for this link
* ``endpoint_a_ref``: Reference to the node at the first end of the link
* ``endpoint_a_port``: The ethernet port or switch port index of the second node
* ``endpoint_b_ref``: Reference to the node at the second end of the link
* ``endpoint_b_port``: The ethernet port or switch port index on the second node
configuration/training_config.rst
configuration/io_settings.rst
configuration/game.rst
configuration/agents.rst
configuration/simulation.rst

174
docs/source/configuration/agents.rst Normal file
View File

@@ -0,0 +1,174 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``agents``
==========
Agents can be scripted (deterministic and stochastic), or controlled by a reinforcement learning algorithm. Not to be confused with an RL agent, the term agent here is used to refer to an entity that sends requests to the simulated network. In this part of the config, each agent's action space, observation space, and reward function can be defined. All three are defined in a modular way.
``agents`` hierarchy
--------------------
.. code-block:: yaml
agents:
- ref: red_agent_example
...
- ref: blue_agent_example
...
- ref: green_agent_example
team: GREEN
type: ProbabilisticAgent
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_2
applications:
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
reward_function:
reward_components:
- type: DUMMY
agent_settings:
start_settings:
start_step: 5
frequency: 4
variance: 3
flatten_obs: False
``ref``
-------
The reference to be used for the given agent.
``team``
--------
Specifies if the agent is malicious (``RED``), benign (``GREEN``), or defensive (``BLUE``). Currently this value is not used for anything other than for human readability in the configuration file.
``type``
--------
Specifies which class should be used for the agent. ``ProxyAgent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``RedDatabaseCorruptingAgent`` and ``ProbabilisticAgent`` generate their own behaviour.
Available agent types:
- ``ProbabilisticAgent``
- ``ProxyAgent``
- ``RedDatabaseCorruptingAgent``
``observation_space``
---------------------
Defines the observation space of the agent.
``type``
^^^^^^^^
selects which python class from the :py:mod:`primaite.game.agent.observation` module is used for the overall observation structure.
``options``
^^^^^^^^^^^
Allows configuration of the chosen observation type. These are optional.
* ``num_services_per_node``, ``num_folders_per_node``, ``num_files_per_folder``, ``num_nics_per_node`` all define the shape of the observation space. The size and shape of the obs space must remain constant, but the number of files, folders, ACL rules, and other components can change within an episode. Therefore padding is performed and these options set the size of the obs space.
* ``nodes``: list of nodes that will be present in this agent's observation space. The ``node_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config. Each node can also be configured with services, and files that should be monitored.
* ``links``: list of links that will be present in this agent's observation space. The ``link_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config.
* ``acl``: configure how the agent reads the access control list on the router in the simulation. ``router_node_ref`` is for selecting which router's ACL table should be used. ``ip_address_order`` sets the encoding of ip addresses as integers within the observation space.
For more information see :py:mod:`primaite.game.agent.observations`
``action_space``
----------------
The action space is configured to be made up of individual action types. Once configured, the agent can select an action type and some optional action parameters at every step. For example: The ``NODE_SERVICE_SCAN`` action takes the parameters ``node_id`` and ``service_id``.
``action_list``
^^^^^^^^^^^^^^^
A list of action modules. The options are listed in the :py:mod:`primaite.game.agent.actions.ActionManager.act_class_identifiers` module.
``action_map``
^^^^^^^^^^^^^^
Restricts the possible combinations of action type / action parameter values to reduce the overall size of the action space. By default, every possible combination of actions and parameters will be assigned an integer for the agent's ``MultiDiscrete`` action space. Instead, the ``action_map`` allows you to list the actions corresponding to each integer in the ``MultiDiscrete`` space.
This is Optional.
``options``
^^^^^^^^^^^
Options that apply to all action components. These are optional.
* ``nodes``: list the nodes that the agent can act on, the order of this list defines the mapping between nodes and ``node_id`` integers.
* ``max_folders_per_node``, ``max_files_per_folder``, ``max_services_per_node``, ``max_nics_per_node``, ``max_acl_rules`` all are used to define the size of the action space.
For more information see :py:mod:`primaite.game.agent.actions`
``reward_function``
-------------------
Similar to action space, this is defined as a list of components from the :py:mod:`primaite.game.agent.rewards` module.
``reward_components``
^^^^^^^^^^^^^^^^^^^^^
A list of reward types from :py:mod:`primaite.game.agent.rewards.RewardFunction.rew_class_identifiers`
e.g.
.. code-block:: yaml
reward_components:
- type: DUMMY
- type: DATABASE_FILE_INTEGRITY
``agent_settings``
------------------
Settings passed to the agent during initialisation. Determines how the agent will behave during training.
e.g.
.. code-block:: yaml
agent_settings:
start_settings:
start_step: 25
frequency: 20
variance: 5
``start_step``
^^^^^^^^^^^^^^
Optional. Default value is ``5``.
The timestep where the agent begins performing actions.
``frequency``
^^^^^^^^^^^^^
Optional. Default value is ``5``.
The number of timesteps the agent will wait before performing another action.
``variance``
^^^^^^^^^^^^
Optional. Default value is ``0``.
The amount of timesteps that the frequency can randomly change.
``flatten_obs``
---------------
If ``True``, gymnasium flattening will be performed on the observation space before sending to the agent. Set this to ``True`` if your agent does not support nested observation spaces.

56
docs/source/configuration/game.rst Normal file
View File

@@ -0,0 +1,56 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``game``
========
This section defines high-level settings that apply across the game, currently it's used to help shape the action and observation spaces by restricting which ports and internet protocols should be considered. Here, users can also set the maximum number of steps in an episode.
``game`` hierarchy
------------------
.. code-block:: yaml
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
thresholds:
nmne:
high: 10
medium: 5
low: 0
``max_episode_length``
----------------------
Optional. Default value is ``256``.
The maximum number of episodes a Reinforcement Learning agent(s) can be trained for.
``ports``
---------
A list of ports that the Reinforcement Learning agent(s) are able to see in the observation space.
See :ref:`List of Ports <List of Ports>` for a list of ports.
``protocols``
-------------
A list of protocols that the Reinforcement Learning agent(s) are able to see in the observation space.
See :ref:`List of IPProtocols <List of IPProtocols>` for a list of protocols.
``thresholds``
--------------
These are used to determine the thresholds of high, medium and low categories for counted observation occurrences.

87
docs/source/configuration/io_settings.rst Normal file
View File

@@ -0,0 +1,87 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``io_settings``
===============
This section configures how PrimAITE saves data during simulation and training.
``io_settings`` hierarchy
-------------------------
.. code-block:: yaml
io_settings:
save_final_model: True
save_checkpoints: False
checkpoint_interval: 10
# save_logs: True
# save_transactions: False
save_agent_actions: True
save_step_metadata: False
save_pcap_logs: False
save_sys_logs: False
``save_final_model``
--------------------
Optional. Default value is ``True``.
Only used if training with PrimaiteSession.
If ``True``, the policy will be saved after the final training iteration.
``save_checkpoints``
--------------------
Optional. Default value is ``False``.
Only used if training with PrimaiteSession.
If ``True``, the policy will be saved periodically during training.
``checkpoint_interval``
-----------------------
Optional. Default value is ``10``.
Only used if training with PrimaiteSession and if ``save_checkpoints`` is ``True``.
Defines how often to save the policy during training.
``save_logs``
-------------
*currently unused*.
``save_agent_actions``
----------------------
Optional. Default value is ``True``.
If ``True``, this will create a JSON file each episode detailing every agent's action in each step of that episode, formatted according to the CAOS format. This includes scripted, RL, and red agents.
``save_step_metadata``
----------------------
Optional. Default value is ``False``.
If ``True``, The RL agent(s) actions, environment states and other data will be saved at every single step.
``save_pcap_logs``
------------------
Optional. Default value is ``False``.
If ``True``, then the pcap files which contain all network traffic during the simulation will be saved.
``save_sys_logs``
-----------------
Optional. Default value is ``False``.
If ``True``, then the log files which contain all node actions during the simulation will be saved.

93
docs/source/configuration/simulation.rst Normal file
View File

@@ -0,0 +1,93 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``simulation``
==============
In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents.
At the top level of the network are ``nodes`` and ``links``.
e.g.
.. code-block:: yaml
simulation:
network:
nodes:
...
links:
...
``nodes``
---------
This is where the list of nodes are defined. Some items will differ according to the node type, however, there will be common items such as a node's reference (which is used by the agent), the node's ``type`` and ``hostname``
To see the configuration for these nodes, refer to the following:
.. toctree::
:maxdepth: 1
:glob:
simulation/nodes/*
``links``
---------
This is where the links between the nodes are formed.
e.g.
In order to recreate the network below, we will need to create 2 links:
- a link from computer_1 to the switch
- a link from computer_2 to the switch
.. image:: ../../_static/switched_p2p_network.png
:width: 500
:align: center
this results in:
.. code-block:: yaml
links:
- ref: computer_1___switch
endpoint_a_ref: computer_1
endpoint_a_port: 1 # port 1 on computer_1
endpoint_b_ref: switch
endpoint_b_port: 1 # port 1 on switch
- ref: computer_2___switch
endpoint_a_ref: computer_2
endpoint_a_port: 1 # port 1 on computer_2
endpoint_b_ref: switch
endpoint_b_port: 2 # port 2 on switch
``ref``
^^^^^^^
The human readable name for the link. Not used in code, however is useful for a human to understand what the link is for.
``endpoint_a_ref``
^^^^^^^^^^^^^^^^^^
The ``hostname`` of the node which must be connected.
``endpoint_a_port``
^^^^^^^^^^^^^^^^^^^
The port on ``endpoint_a_ref`` which is to be connected to ``endpoint_b_port``.
This accepts an integer value e.g. if port 1 is to be connected, the configuration should be ``endpoint_a_port: 1``
``endpoint_b_ref``
^^^^^^^^^^^^^^^^^^
The ``hostname`` of the node which must be connected.
``endpoint_b_port``
^^^^^^^^^^^^^^^^^^^
The port on ``endpoint_b_ref`` which is to be connected to ``endpoint_a_port``.
This accepts an integer value e.g. if port 1 is to be connected, the configuration should be ``endpoint_b_port: 1``

35
docs/source/configuration/simulation/nodes/common/common.rst Normal file
View File

@@ -0,0 +1,35 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _Node Attributes:
Common Attributes
#################
Node Attributes
===============
Attributes that are shared by all nodes.
.. include:: common_node_attributes.rst
.. _Network Node Attributes:
Network Node Attributes
=======================
Attributes that are shared by nodes that inherit from :py:mod:`primaite.simulator.network.hardware.nodes.network.network_node.NetworkNode`
.. include:: common_host_node_attributes.rst
.. _Host Node Attributes:
Host Node Attributes
====================
Attributes that are shared by nodes that inherit from :py:mod:`primaite.simulator.network.hardware.nodes.host.host_node.HostNode`
.. include:: common_host_node_attributes.rst
.. |NODE| replace:: node

26
docs/source/configuration/simulation/nodes/common/common_host_node_attributes.rst Normal file
View File

@@ -0,0 +1,26 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _common_host_node_attributes:
``ip_address``
--------------
The IP address of the |NODE| in the network.
``subnet_mask``
---------------
Optional. Default value is ``255.255.255.0``.
The subnet mask for the |NODE| to use.
``default_gateway``
-------------------
The IP address that the |NODE| will use as the default gateway. Typically, this is the IP address of the closest router that the |NODE| is connected to.
.. include:: ../software/applications.rst
.. include:: ../software/services.rst

51
docs/source/configuration/simulation/nodes/common/common_network_node_attributes.rst Normal file
View File

@@ -0,0 +1,51 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _common_network_node_attributes:
``routes``
----------
A list of routes which tells the |NODE| where to forward the packet to depending on the target IP address.
e.g.
.. code-block:: yaml
nodes:
- ref: node
...
routes:
- address: 192.168.0.10
subnet_mask: 255.255.255.0
next_hop_ip_address: 192.168.1.1
metric: 0
``address``
"""""""""""
The target IP address for the route. If the packet destination IP address matches this, the |NODE| will route the packet according to the ``next_hop_ip_address``.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``subnet_mask``
"""""""""""""""
Optional. Default value is ``255.255.255.0``.
The subnet mask setting for the route.
``next_hop_ip_address``
"""""""""""""""""""""""
The IP address of the next hop IP address that the packet will follow if the address matches the packet's destination IP address.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``metric``
""""""""""
Optional. Default value is ``0``. This value accepts floats.
The cost or distance of a route. The higher the value, the more cost or distance is attributed to the route.

55
docs/source/configuration/simulation/nodes/common/common_node_attributes.rst Normal file
View File

@@ -0,0 +1,55 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _common_node_attributes:
``ref``
-------
Human readable name used as reference for the |NODE|. Not used in code.
``hostname``
------------
The hostname of the |NODE|. This will be used to reference the |NODE|.
``operating_state``
-------------------
The initial operating state of the node.
Optional. Default value is ``ON``.
Options available are:
- ``ON``
- ``OFF``
- ``BOOTING``
- ``SHUTTING_DOWN``
Note that YAML may assume non quoted ``ON`` and ``OFF`` as ``True`` and ``False`` respectively. To prevent this, use ``"ON"`` or ``"OFF"``
See :py:mod:`primaite.simulator.network.hardware.node_operating_state.NodeOperatingState`
``dns_server``
--------------
Optional. Default value is ``None``.
The IP address of the node which holds an instance of the :ref:`DNSServer`. Some applications may use a domain name e.g. the :ref:`WebBrowser`
``start_up_duration``
---------------------
Optional. Default value is ``3``.
The number of time steps required to occur in order for the node to cycle from ``OFF`` to ``BOOTING_UP`` and then finally ``ON``.
``shut_down_duration``
----------------------
Optional. Default value is ``3``.
The number of time steps required to occur in order for the node to cycle from ``ON`` to ``SHUTTING_DOWN`` and then finally ``OFF``.

18
docs/source/configuration/simulation/nodes/common/node_type_list.rst Normal file
View File

@@ -0,0 +1,18 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``type``
--------
The type of node to add.
Available options are:
- ``computer``
- ``firewall``
- ``router``
- ``server``
- ``switch``
To create a |NODE|, type must be |NODE_TYPE|.

41
docs/source/configuration/simulation/nodes/computer.rst Normal file
View File

@@ -0,0 +1,41 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _computer_configuration:
``computer``
============
A basic representation of a computer within the simulation.
See :py:mod:`primaite.simulator.network.hardware.nodes.host.computer.Computer`
example computer
----------------
.. code-block:: yaml
simulation:
network:
nodes:
- ref: client_1
hostname: client_1
type: computer
ip_address: 192.168.0.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.0.1
dns_server: 192.168.1.10
applications:
...
services:
...
.. include:: common/common_node_attributes.rst
.. include:: common/node_type_list.rst
.. include:: common/common_host_node_attributes.rst
.. |NODE| replace:: computer
.. |NODE_TYPE| replace:: ``computer``

300
docs/source/configuration/simulation/nodes/firewall.rst Normal file
View File

@@ -0,0 +1,300 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _firewall_configuration:
``firewall``
============
A basic representation of a network firewall within the simulation.
The firewall is similar to how :ref:`Router <router_configuration>` works, with the difference being how firewall has specific ACL rules for inbound and outbound traffic as well as firewall being limited to 3 ports.
See :py:mod:`primaite.simulator.network.hardware.nodes.network.firewall.Firewall`
example firewall
----------------
.. code-block:: yaml
simulation:
network:
nodes:
- ref: firewall
hostname: firewall
type: firewall
start_up_duration: 0
shut_down_duration: 0
ports:
external_port: # port 1
ip_address: 192.168.20.1
subnet_mask: 255.255.255.0
internal_port: # port 2
ip_address: 192.168.1.2
subnet_mask: 255.255.255.0
dmz_port: # port 3
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
acl:
internal_inbound_acl:
...
internal_outbound_acl:
...
dmz_inbound_acl:
...
dmz_outbound_acl:
...
external_inbound_acl:
...
external_outbound_acl:
...
routes:
...
.. include:: common/common_node_attributes.rst
.. include:: common/node_type_list.rst
``ports``
---------
The firewall node only has 3 ports. These specifically are:
- ``external_port`` (port 1)
- ``internal_port`` (port 2)
- ``dmz_port`` (port 3) (can be optional)
The ports should be defined with an ip address and subnet mask e.g.
.. code-block:: yaml
nodes:
- ref: firewall
...
ports:
external_port: # port 1
ip_address: 192.168.20.1
subnet_mask: 255.255.255.0
internal_port: # port 2
ip_address: 192.168.1.2
subnet_mask: 255.255.255.0
dmz_port: # port 3
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
``ip_address``
""""""""""""""
The IP address for the given port. This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``subnet_mask``
"""""""""""""""
Optional. Default value is ``255.255.255.0``.
The subnet mask setting for the port.
``acl``
-------
There are 6 ACLs that can be defined for a firewall
- ``internal_inbound_acl`` for traffic going towards the internal network
- ``internal_outbound_acl`` for traffic coming from the internal network
- ``dmz_inbound_acl`` for traffic going towards the dmz network
- ``dmz_outbound_acl`` for traffic coming from the dmz network
- ``external_inbound_acl`` for traffic coming from the external network
- ``external_outbound_acl`` for traffic going towards the external network
.. image:: ../../../../_static/firewall_acl.png
:width: 500
:align: center
By default, ``external_inbound_acl`` and ``external_outbound_acl`` will permit any traffic through.
``internal_inbound_acl``, ``internal_outbound_acl``, ``dmz_inbound_acl`` and ``dmz_outbound_acl`` will deny any traffic by default, so must be configured to allow defined ``src_port`` and ``dst_port`` or ``protocol``.
See :py:mod:`primaite.simulator.network.hardware.nodes.network.router.AccessControlList`
See :ref:`List of Ports <List of Ports>` for a list of ports.
``internal_inbound_acl``
""""""""""""""""""""""""
ACL rules for packets that have a destination IP address in what is considered the internal network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
internal_inbound_acl:
21: # position 21 on ACL list
action: PERMIT # allow packets that
src_port: POSTGRES_SERVER # are emitted from the POSTGRES_SERVER port
dst_port: POSTGRES_SERVER # are going towards an POSTGRES_SERVER port
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
``internal_outbound_acl``
"""""""""""""""""""""""""
ACL rules for packets that have a source IP address in what is considered the internal network and is going towards the DMZ network or the external network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
internal_outbound_acl:
21: # position 21 on ACL list
action: PERMIT # allow packets that
src_port: POSTGRES_SERVER # are emitted from the POSTGRES_SERVER port
dst_port: POSTGRES_SERVER # are going towards an POSTGRES_SERVER port
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
``dmz_inbound_acl``
"""""""""""""""""""
ACL rules for packets that have a destination IP address in what is considered the DMZ network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
dmz_inbound_acl:
19: # position 19 on ACL list
action: PERMIT # allow packets that
src_port: POSTGRES_SERVER # are emitted from the POSTGRES_SERVER port
dst_port: POSTGRES_SERVER # are going towards an POSTGRES_SERVER port
20: # position 20 on ACL list
action: PERMIT # allow packets that
src_port: HTTP # are emitted from the HTTP port
dst_port: HTTP # are going towards an HTTP port
21: # position 21 on ACL list
action: PERMIT # allow packets that
src_port: HTTPS # are emitted from the HTTPS port
dst_port: HTTPS # are going towards an HTTPS port
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
``dmz_outbound_acl``
""""""""""""""""""""
ACL rules for packets that have a source IP address in what is considered the DMZ network and is going towards the internal network or the external network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
dmz_outbound_acl:
19: # position 19 on ACL list
action: PERMIT # allow packets that
src_port: POSTGRES_SERVER # are emitted from the POSTGRES_SERVER port
dst_port: POSTGRES_SERVER # are going towards an POSTGRES_SERVER port
20: # position 20 on ACL list
action: PERMIT # allow packets that
src_port: HTTP # are emitted from the HTTP port
dst_port: HTTP # are going towards an HTTP port
21: # position 21 on ACL list
action: PERMIT # allow packets that
src_port: HTTPS # are emitted from the HTTPS port
dst_port: HTTPS # are going towards an HTTPS port
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
``external_inbound_acl``
""""""""""""""""""""""""
Optional. By default, this will allow any traffic through.
ACL rules for packets that have a destination IP address in what is considered the external network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
external_inbound_acl:
21: # position 19 on ACL list
action: DENY # deny packets that
src_port: POSTGRES_SERVER # are emitted from the POSTGRES_SERVER port
dst_port: POSTGRES_SERVER # are going towards an POSTGRES_SERVER port
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
``external_outbound_acl``
"""""""""""""""""""""""""
Optional. By default, this will allow any traffic through.
ACL rules for packets that have a source IP address in what is considered the external network and is going towards the DMZ network or the internal network.
example:
.. code-block:: yaml
nodes:
- ref: firewall
...
acl:
external_outbound_acl:
22: # position 22 on ACL list
action: PERMIT # allow packets that
src_port: ARP # are emitted from the ARP port
dst_port: ARP # are going towards an ARP port
23: # position 23 on ACL list
action: PERMIT # allow packets that
protocol: ICMP # are ICMP
.. include:: common/common_network_node_attributes.rst
.. |NODE| replace:: firewall
.. |NODE_TYPE| replace:: ``firewall``

127
docs/source/configuration/simulation/nodes/router.rst Normal file
View File

@@ -0,0 +1,127 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _router_configuration:
``router``
==========
A basic representation of a network router within the simulation.
See :py:mod:`primaite.simulator.network.hardware.nodes.network.router.Router`
example router
--------------
.. code-block:: yaml
simulation:
network:
nodes:
- ref: router_1
hostname: router_1
type: router
num_ports: 5
ports:
...
acl:
...
.. include:: common/common_node_attributes.rst
.. include:: common/node_type_list.rst
``num_ports``
-------------
Optional. Default value is ``5``.
The number of ports the router will have.
``ports``
---------
Sets up the router's ports with an IP address and a subnet mask.
Example of setting ports for a router with 2 ports:
.. code-block:: yaml
nodes:
- ref: router_1
...
ports:
1:
ip_address: 192.168.1.1
subnet_mask: 255.255.255.0
2:
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
``ip_address``
""""""""""""""
The IP address for the given port. This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``subnet_mask``
"""""""""""""""
Optional. Default value is ``255.255.255.0``.
The subnet mask setting for the port.
``acl``
-------
Sets up the ACL rules for the router.
e.g.
.. code-block:: yaml
nodes:
- ref: router_1
...
acl:
1:
action: PERMIT
src_port: ARP
dst_port: ARP
2:
action: PERMIT
protocol: ICMP
See :py:mod:`primaite.simulator.network.hardware.nodes.network.router.AccessControlList`
See :ref:`List of Ports <List of Ports>` for a list of ports.
``action``
""""""""""
Available options are
- ``PERMIT`` : Allows the specified ``protocol`` or ``src_port`` and ``dst_port`` pairs
- ``DENY`` : Blocks the specified ``protocol`` or ``src_port`` and ``dst_port`` pairs
``src_port``
""""""""""""
Is used alongside ``dst_port``. Specifies the port where a packet originates. Used by the ACL Rule to determine if a packet with a specific source port is allowed to pass through the network node.
``dst_port``
""""""""""""
Is used alongside ``src_port``. Specifies the port where a packet is destined to arrive. Used by the ACL Rule to determine if a packet with a specific destination port is allowed to pass through the network node.
``protocol``
""""""""""""
Specifies which protocols are allowed by the ACL Rule to pass through the network node.
See :ref:`List of IPProtocols <List of IPProtocols>` for a list of protocols.
.. include:: common/common_network_node_attributes.rst
.. |NODE| replace:: router
.. |NODE_TYPE| replace:: ``router``

41
docs/source/configuration/simulation/nodes/server.rst Normal file
View File

@@ -0,0 +1,41 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _server_configuration:
``server``
==========
A basic representation of a server within the simulation.
See :py:mod:`primaite.simulator.network.hardware.nodes.host.server.Server`
example server
--------------
.. code-block:: yaml
simulation:
network:
nodes:
- ref: server_1
hostname: server_1
type: server
ip_address: 192.168.10.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
...
services:
...
.. include:: common/common_node_attributes.rst
.. include:: common/node_type_list.rst
.. include:: common/common_host_node_attributes.rst
.. |NODE| replace:: server
.. |NODE_TYPE| replace:: ``server``

39
docs/source/configuration/simulation/nodes/switch.rst Normal file
View File

@@ -0,0 +1,39 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _switch_configuration:
``switch``
==========
A basic representation of a network switch within the simulation.
See :py:mod:`primaite.simulator.network.hardware.nodes.network.switch.Switch`
example switch
--------------
.. code-block:: yaml
simulation:
network:
nodes:
- ref: switch_1
hostname: switch_1
type: switch
num_ports: 8
.. include:: common/common_node_attributes.rst
.. include:: common/node_type_list.rst
``num_ports``
-------------
Optional. Default value is ``8``.
The number of ports the switch will have.
.. |NODE| replace:: switch
.. |NODE_TYPE| replace:: ``switch``

25
docs/source/configuration/simulation/software/applications.rst Normal file
View File

@@ -0,0 +1,25 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``applications``
----------------
List of available applications that can be installed on a |NODE| can be found in :ref:`List of Applications <List of Applications>`
application in configuration
""""""""""""""""""""""""""""
Applications takes a list of applications as shown in the example below.
.. code-block:: yaml
- ref: client_1
hostname: client_1
type: computer
...
applications:
- ref: example_application
type: example_application_type
options:
# this section is different for each application

25
docs/source/configuration/simulation/software/services.rst Normal file
View File

@@ -0,0 +1,25 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``services``
------------
List of available services that can be installed on a |NODE| can be found in :ref:`List of Services <List of Services>`
services in configuration
"""""""""""""""""""""""""
Services takes a list of services as shown in the example below.
.. code-block:: yaml
- ref: client_1
hostname: client_1
type: computer
...
applications:
- ref: example_service
type: example_service_type
options:
# this section is different for each service

75
docs/source/configuration/training_config.rst Normal file
View File

@@ -0,0 +1,75 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``training_config``
===================
Configuration items relevant to how the Reinforcement Learning agent(s) will be trained.
``training_config`` hierarchy
-----------------------------
.. code-block:: yaml
training_config:
rl_framework: SB3 # or RLLIB_single_agent or RLLIB_multi_agent
rl_algorithm: PPO # or A2C
n_learn_episodes: 5
max_steps_per_episode: 200
n_eval_episodes: 1
deterministic_eval: True
seed: 123
``rl_framework``
----------------
The RL (Reinforcement Learning) Framework to use in the training session
Options available are:
- ``SB3`` (Stable Baselines 3)
- ``RLLIB_single_agent`` (Single Agent Ray RLLib)
- ``RLLIB_multi_agent`` (Multi Agent Ray RLLib)
``rl_algorithm``
----------------
The Reinforcement Learning Algorithm to use in the training session
Options available are:
- ``PPO`` (Proximal Policy Optimisation)
- ``A2C`` (Advantage Actor Critic)
``n_learn_episodes``
--------------------
The number of episodes to train the agent(s).
This should be an integer value above ``0``
``max_steps_per_episode``
-------------------------
The number of steps each episode will last for.
This should be an integer value above ``0``.
``n_eval_episodes``
-------------------
Optional. Default value is ``0``.
The number of evaluation episodes to run the trained agent for.
This should be an integer value above ``0``.
``deterministic_eval``
----------------------
Optional. By default this value is ``False``.
If this is set to ``True``, the agents will act deterministically instead of stochastically.
``seed``
--------
Optional.
The seed is used (alongside ``deterministic_eval``) to reproduce a previous instance of training and evaluation of an RL agent.
The seed should be an integer value.
Useful for debugging.

4
docs/source/customising_scenarios.rst Normal file
View File

@@ -0,0 +1,4 @@
Customising Agents
******************
For an example of how to customise red agent behaviour in the Data Manipulation scenario, please refer to the notebook ``Data-Manipulation-Customising-Red-Agent.ipynb``.

2
docs/source/dependencies.rst
View File

@@ -5,6 +5,8 @@
.. role:: raw-html(raw)
:format: html
.. _Dependencies:
Dependencies
============

10
docs/source/environment.rst Normal file
View File

@@ -0,0 +1,10 @@
RL Environments
***************
RL environments are the objects that directly interface with RL libraries such as Stable-Baselines3 and Ray RLLib. The PrimAITE simulation is exposed via three different environment APIs:
* Gymnasium API - this is the standard interface that works with many RL libraries like SB3, Ray, Tianshou, etc. ``PrimaiteGymEnv`` adheres to the `Official Gymnasium documentation <https://gymnasium.farama.org/api/env/>`_.
* Ray Single agent API - For training a single Ray RLLib agent
* Ray MARL API - For training multi-agent systems with Ray RLLib. ``PrimaiteRayMARLEnv`` adheres to the `Official Ray documentation <https://docs.ray.io/en/latest/rllib/package_ref/env/multi_agent_env.html>`_.
There are Jupyter notebooks which demonstrate integration with each of these three environments. They are located in ``~/primaite/<VERSION>/notebooks/example_notebooks``.

77
docs/source/example_notebooks.rst Normal file
View File

@@ -0,0 +1,77 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Example Jupyter Notebooks
=========================
There are a few example notebooks included which help with the understanding of PrimAITE's capabilities.
The Jupyter Notebooks can be run via the 2 examples below. These assume that the instructions to install PrimAITE from the :ref:`Getting Started <getting-started>` page is completed as a prerequisite.
Running Jupyter Notebooks
-------------------------
1. Navigate to the PrimAITE directory
.. code-block:: bash
:caption: Unix
cd ~/primaite/{VERSION}
.. code-block:: powershell
:caption: Windows (Powershell)
cd ~\primaite\{VERSION}
2. Run jupyter notebook (the python environment to which you installed PrimAITE must be active)
.. code-block:: bash
:caption: Unix
jupyter notebook
.. code-block:: powershell
:caption: Windows (Powershell)
jupyter notebook
3. Opening the jupyter webpage (optional)
The default web browser may automatically open the webpage. However, if that is not the case, click the link shown in your command prompt output. It should look like this: ``http://localhost:8888/?token=0123456798abc0123456789abc``
4. Navigate to the list of notebooks
The example notebooks are located in ``notebooks/example_notebooks/``. The file system shown in the jupyter webpage is relative to the location in which the ``jupyter notebook`` command was used.
Running Jupyter Notebooks via VSCode
------------------------------------
It is also possible to view the Jupyter notebooks within VSCode.
The best place to start is by opening a notebook file (.ipynb) in VSCode. If using VSCode to view a notebook for the first time, follow the steps below.
Installing extensions
"""""""""""""""""""""
VSCode may need some extensions to be installed if not already done.
To do this, press the "Select Kernel" button on the top right.
This should open a dialog which has the option to install python and jupyter extensions.
.. image:: ../../_static/notebooks/install_extensions.png
:width: 700
:align: center
:alt: :: The top dialog option that appears will automatically install the extensions
The following extensions should now be installed
.. image:: ../../_static/notebooks/extensions.png
:width: 300
:align: center
VSCode will then ask for a Python environment version to use. PrimAITE is compatible with Python versions 3.8 - 3.10
You should now be able to interact with the notebook.

81
docs/source/game_layer.rst
View File

@@ -6,44 +6,91 @@ The Primaite codebase consists of two main modules:
* ``simulator``: The simulation logic including the network topology, the network state, and behaviour of various hardware and software classes.
* ``game``: The agent-training infrastructure which helps reinforcement learning agents interface with the simulation. This includes the observation, action, and rewards, for RL agents, but also scripted deterministic agents. The game layer orchestrates all the interactions between modules.
The simulator and game layer communicate using the PrimAITE State API and the PrimAITE Request API.
..
TODO: write up these APIs and link them here.
Game layer
----------
The simulator and game layer communicate using the PrimAITE State API and the PrimAITE Request API.
The game layer is responsible for managing agents and getting them to interface with the simulator correctly. It consists of several components:
PrimAITE Session
^^^^^^^^^^^^^^^
================
.. admonition:: Deprecated
:class: deprecated
PrimAITE Session is being deprecated in favour of Jupyter Notebooks. The `session` command will be removed in future releases, but example notebooks will be provided to demonstrate the same functionality.
``PrimaiteSession`` is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. ``PrimaiteSession`` keeps track of multiple agents of different types.
Agents
^^^^^^
======
All agents inherit from the :py:class:`primaite.game.agent.interface.AbstractAgent` class, which mandates that they have an ObservationManager, ActionManager, and RewardManager. The agent behaviour depends on the type of agent, but there are two main types:
* RL agents action during each step is decided by an appropriate RL algorithm. The agent within PrimAITE just acts to format and forward actions decided by an RL policy.
* Deterministic agents perform all of their decision making within the PrimAITE game layer. They typically have a scripted policy which always performs the same action or a rule-based policy which performs actions based on the current state of the simulation. They can have a stochastic element, and their seed will be settable.
* Deterministic agents perform all of their decision making within the PrimAITE game layer. They typically have a scripted policy which always performs the same action or a rule-based policy which performs actions based on the current state of the simulation. They can have a stochastic element, and their seed is settable.
..
TODO: add seed to stochastic scripted agents
Observations
^^^^^^^^^^^^^^^^^^
============
An agent's observations are managed by the ``ObservationManager`` class. It generates observations based on the current simulation state dictionary. It also provides the observation space during initial setup. The data is formatted so it's compatible with ``Gymnasium.spaces``. Observation spaces are composed of one or more components which are defined by the ``AbstractObservation`` base class.
Actions
^^^^^^^
=======
An agent's actions are managed by the ``ActionManager``. It converts actions selected by agents (which are typically integers chosen from a ``gymnasium.spaces.Discrete`` space) into simulation-friendly requests. It also provides the action space during initial setup. Action spaces are composed of one or more components which are defined by the ``AbstractAction`` base class.
Rewards
^^^^^^^
=======
An agent's reward function is managed by the ``RewardManager``. It calculates rewards based on the simulation state (in a way similar to observations). Rewards can be defined as a weighted sum of small reward components. For example, an agents reward can be based on the uptime of a database service plus the loss rate of packets between clients and a web server. The reward components are defined by the AbstractReward base class.
An agent's reward function is managed by the ``RewardManager``. It calculates rewards based on the simulation state (in a way similar to observations). Rewards can be defined as a weighted sum of small reward components. For example, an agents reward can be based on the uptime of a database service plus the loss rate of packets between clients and a web server.
Reward Components
-----------------
Currently implemented are reward components tailored to the data manipulation scenario. View the full API and description of how they work here: :py:module:`primaite.game.agent.reward`.
Reward Sharing
--------------
An agent's reward can be based on rewards of other agents. This is particularly useful for modelling a situation where the blue agent's job is to protect the ability of green agents to perform their pattern-of-life. This can be configured in the YAML file this way:
```yaml
green_agent_1: # this agent sometimes tries to access the webpage, and sometimes the database
# actions, observations, and agent settings go here
reward_function:
reward_components:
# When the webpage loads, the reward goes up by 0.25 when it fails to load, it goes down to -0.25
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.25
options:
node_hostname: client_2
# When the database is reachable, the reward goes up by 0.05, when it is unreachable it goes down to -0.05
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
weight: 0.05
options:
node_hostname: client_2
blue_agent:
# actions, observations, and agent settings go here
reward_function:
reward_components:
# When the database file is in a good state, blue's reward is 0.4, when it's in a corrupted state the reward is -0.4
- type: DATABASE_FILE_INTEGRITY
weight: 0.40
options:
node_hostname: database_server
folder_name: database
file_name: database.db
# The green's reward is added onto the blue's reward.
- type: SHARED_REWARD
weight: 1.0
options:
agent_name: client_2_green_user
```
When defining agent reward sharing, users must be careful to avoid circular references, as that would lead to an infinite calculation loop. PrimAITE will prevent circular dependencies and provide a helpful error message if they are detected in the yaml.

12
docs/source/getting_started.rst
View File

@@ -11,7 +11,7 @@ Getting Started
Pre-Requisites
In order to get **PrimAITE** installed, you will need to have a python version between 3.8 and 3.11 installed. If you don't already have it, this is how to install it:
In order to get **PrimAITE** installed, you will need Python, venv, and pip. If you don't already have them, this is how to install it:
.. code-block:: bash
@@ -30,6 +30,8 @@ In order to get **PrimAITE** installed, you will need to have a python version b
**PrimAITE** is designed to be OS-agnostic, and thus should work on most variations/distros of Linux, Windows, and MacOS.
Installing PrimAITE has been tested with all supported python versions, venv 20.24.1, and pip 23.
Install PrimAITE
****************
@@ -38,12 +40,12 @@ Install PrimAITE
.. code-block:: bash
:caption: Unix
mkdir ~/primaite/3.0.0
mkdir -p ~/primaite/{VERSION}
.. code-block:: powershell
:caption: Windows (Powershell)
mkdir ~\primaite\3.0.0
mkdir ~\primaite\{VERSION}
2. Navigate to the primaite directory and create a new python virtual environment (venv)
@@ -51,13 +53,13 @@ Install PrimAITE
.. code-block:: bash
:caption: Unix
cd ~/primaite/3.0.0
cd ~/primaite/{VERSION}
python3 -m venv .venv
.. code-block:: powershell
:caption: Windows (Powershell)
cd ~\primaite\3.0.0
cd ~\primaite\{VERSION}
python3 -m venv .venv
attrib +h .venv /s /d # Hides the .venv directory

9
docs/source/primaite_session.rst
View File

@@ -4,6 +4,11 @@
.. _run a primaite session:
.. admonition:: Deprecated
:class: deprecated
PrimAITE Session is being deprecated in favour of Jupyter Notebooks. The ``session`` command will be removed in future releases, but example notebooks will be provided to demonstrate the same functionality.
Run a PrimAITE Session
======================
@@ -30,7 +35,7 @@ Outputs
-------
Running a session creates a session output directory in your user data folder. The filepath looks like this:
``~/primaite/3.0.0/sessions/YYYY-MM-DD/HH-MM-SS/``. This folder contains the simulation sys logs generated by each node,
``~/primaite/{VERSION}/sessions/YYYY-MM-DD/HH-MM-SS/``. This folder contains the simulation sys logs generated by each node,
the saved agent checkpoints, and final model. The folder also contains a .json file for each episode step that
contains the action, reward, and simulation state. These can be found in
``~/primaite/3.0.0/sessions/YYYY-MM-DD/HH-MM-SS/simulation_output/episode_<n>/step_metadata/step_<n>.json``
``~/primaite/{VERSION}/sessions/YYYY-MM-DD/HH-MM-SS/simulation_output/episode_<n>/step_metadata/step_<n>.json``

45
docs/source/request_system.rst
View File

@@ -3,7 +3,7 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Request System
==============
**************
``SimComponent`` objects in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``.
@@ -12,31 +12,42 @@ Just like other aspects of SimComponent, the request types are not managed centr
- API
When requesting an action within the simulation, these two arguments must be provided:
1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '<node-name>', 'service', '<service-name>', 'restart']`.
1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as ``['network', 'node', '<node-name>', 'service', '<service-name>', 'restart']``.
2. ``context`` - optional extra information that can be used to decide how to process the request. This is formatted as a dictionary. For example, if the request requires authentication, the context can include information about the user that initiated the request to decide if their permissions are sufficient.
When a request is resolved, it returns a success status, and optional additional data about the request.
``status`` can be one of:
* ``success``: the request was executed
* ``failure``: the request could not be executed
* ``unreachable``: the target for the request was not found
* ``pending``: the request was initiated, but has not finished during this step
``data`` can be a dictionary with any arbitrary JSON-like data to describe the outcome of the request.
- ``request`` detail
The request is a list of strings which help specify who should handle the request. The strings in the request list help RequestManagers traverse the 'ownership tree' of SimComponent. The example given above would be handled in the following way:
1. ``Simulation`` receives `['network', 'node', '<node-name>', 'service', '<service-name>', 'restart']`.
1. ``Simulation`` receives ``['network', 'node', 'computer_1', 'service', 'DNSService', 'restart']``.
The first element of the request is ``network``, therefore it passes the request down to its network.
2. ``Network`` receives `['node', '<node-name>', 'service', '<service-name>', 'restart']`.
2. ``Network`` receives ``['node', 'computer_1', 'service', 'DNSService', 'restart']``.
The first element of the request is ``node``, therefore the network looks at the node name and passes the request down to the node with that name.
3. ``Node`` receives `['service', '<service-name>', 'restart']`.
3. ``computer_1`` receives ``['service', 'DNSService', 'restart']``.
The first element of the request is ``service``, therefore the node looks at the service name and passes the rest of the request to the service with that name.
4. ``Service`` receives ``['restart']``.
4. ``DNSService`` receives ``['restart']``.
Since ``restart`` is a defined request type in the service's own RequestManager, the service performs a restart.
- ``context`` detail
The context is not used by any of the currently implemented components or requests.
Technical Detail
----------------
================
This system was achieved by implementing two classes, :py:class:`primaite.simulator.core.RequestType`, and :py:class:`primaite.simulator.core.RequestManager`.
``RequestType``
------
---------------
The ``RequestType`` object stores a reference to a method that executes the request, for example a node could have a request type that stores a reference to ``self.turn_on()``. Technically, this can be any callable that accepts `request, context` as it's parameters. In practice, this is often defined using ``lambda`` functions within a component's ``self._init_request_manager()`` method. Optionally, the ``RequestType`` object can also hold a validator that will permit/deny the request depending on context.
@@ -60,7 +71,7 @@ A simple example without chaining can be seen in the :py:class:`primaite.simulat
*ellipses (``...``) used to omit code impertinent to this explanation*
Chaining RequestManagers
-----------------------
------------------------
A request function needs to be a callable that accepts ``request, context`` as parameters. Since the request manager resolves requests by invoking it with ``request, context`` as parameter, it is possible to use a ``RequestManager`` as a ``RequestType``.
@@ -93,3 +104,19 @@ An example of how this works is in the :py:class:`primaite.simulator.network.har
self._service_request_manager.add_request(service.name, RequestType(func=service._request_manager))
This process is repeated until the request word corresponds to a callable function rather than another ``RequestManager`` .
Request Validation
------------------
There are times when a request should be rejected. For instance, if an agent attempts to run an application on a node that is currently off. For this purpose, requests are filtered by an object called a validator. :py:class:`primaite.simulator.core.RequestPermissionValidator` is a basic class whose ``__call__()`` method returns ``True`` if the request should be permitted or ``False`` if it cannot be permitted. For example, the Node class has a validator called :py:class:`primaite.simulator.network.hardware.base.Node._NodeIsOnValidator<_NodeIsOnValidator>` which allows requests only when the operating status of the node is ``ON``.
Requests that are specified without a validator automatically get assigned an ``AllowAllValidator`` which allows requests no matter what.
Request Response
----------------
The :py:class:`primaite.interface.request.RequestResponse<RequestResponse>` is a data transfer object that carries response data between the simulator and the game layer. The ``status`` field reports on the success or failure, and the ``data`` field is for any additional data. The most common way that this class is initiated is by its ``from_bool`` method. This way, given a True or False, a successful or failed request response is generated, respectively (with an empty data field).
For instance, the ``execute`` action on a :py:class:`primaite.simulator.system.applications.web_browser.WebBrowser<WebBrowser>` calls the ``get_webpage()`` method of the ``WebBrowser``. ``get_webpage()`` returns a True if the webpage was successfully retrieved, and False if unsuccessful for any reason, such as being blocked by an ACL, or if the database server is unresponsive. The boolean returned from ``get_webpage()`` is used to create the request response.
Just as the requests themselves were passed from owner to component, the request response is bubbled back up from component to owner until it arrives at the game layer.

2
docs/source/simulation.rst
View File

@@ -22,9 +22,9 @@ Contents
simulation_components/network/nodes/host_node
simulation_components/network/nodes/network_node
simulation_components/network/nodes/router
simulation_components/network/nodes/switch
simulation_components/network/nodes/wireless_router
simulation_components/network/nodes/firewall
simulation_components/network/switch
simulation_components/network/network
simulation_components/system/internal_frame_processing
simulation_components/system/sys_log

79
docs/source/simulation_components/network/base_hardware.rst
View File

@@ -12,34 +12,81 @@ complex, specialized hardware components inherit from and build upon.
The key elements defined in ``base.py`` are:
NetworkInterface
================
``NetworkInterface``
====================
- Abstract base class for network interfaces like NICs. Defines common attributes like MAC address, speed, MTU.
- Requires subclasses to implement ``enable()``, ``disable()``, ``send_frame()`` and ``receive_frame()``.
- Provides basic state description and request handling capabilities.
Node
====
``Node``
========
The Node class stands as a central component in ``base.py``, acting as the superclass for all network nodes within a
PrimAITE simulation.
Node Attributes
---------------
See :ref:`Node Attributes`
- **hostname**: The network hostname of the node.
- **operating_state**: Indicates the current hardware state of the node.
- **network_interfaces**: Maps interface names to NetworkInterface objects on the node.
- **network_interface**: Maps port IDs to ``NetworkInterface`` objects on the node.
- **dns_server**: Specifies DNS servers for domain name resolution.
- **start_up_duration**: The time it takes for the node to become fully operational after being powered on.
- **shut_down_duration**: The time required for the node to properly shut down.
- **sys_log**: A system log for recording events related to the node.
- **session_manager**: Manages user sessions within the node.
- **software_manager**: Controls the installation and management of software and services on the node.
.. _Node Start up and Shut down:
Node Start up and Shut down
---------------------------
Nodes are powered on and off over multiple timesteps. By default, the node ``start_up_duration`` and ``shut_down_duration`` is 3 timesteps.
Example code where a node is turned on:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a")
assert node.operating_state is NodeOperatingState.OFF # By default, node is instantiated in an OFF state
node.power_on() # power on the node
assert node.operating_state is NodeOperatingState.BOOTING # node is booting up
for i in range(node.start_up_duration + 1):
# apply timestep until the node start up duration
node.apply_timestep(timestep=i)
assert node.operating_state is NodeOperatingState.ON # node is in ON state
If the node needs to be instantiated in an on state:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a", operating_state=NodeOperatingState.ON)
assert node.operating_state is NodeOperatingState.ON # node is in ON state
Setting ``start_up_duration`` and/or ``shut_down_duration`` to ``0`` will allow for the ``power_on`` and ``power_off`` methods to be completed instantly without applying timesteps:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a", start_up_duration=0, shut_down_duration=0)
assert node.operating_state is NodeOperatingState.OFF # node is in OFF state
node.power_on()
assert node.operating_state is NodeOperatingState.ON # node is in ON state
node.power_off()
assert node.operating_state is NodeOperatingState.OFF # node is in OFF state
Node Behaviours/Functions
-------------------------

10
docs/source/simulation_components/network/network.rst
View File

@@ -30,11 +30,11 @@ we'll use the following Network that has a client, server, two switches, and a r
.. code-block:: python
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import NIC
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import Router, ACLAction
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.network.router import Router, ACLAction
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port

13
docs/source/simulation_components/network/network_interfaces.rst
View File

@@ -13,6 +13,8 @@ facilitates modular development, enhances maintainability, and supports scalabil
allowing for focused enhancements within each layer.
.. image:: primaite_network_interface_model.png
:width: 500
:align: center
Layer Descriptions
==================
@@ -65,9 +67,14 @@ Network Interface Classes
**NetworkInterface (Base Layer)**
Abstract base class defining core interface properties like MAC address, speed, MTU.
Requires subclasses implement key methods like send/receive frames, enable/disable interface.
Establishes universal network interface capabilities.
- Abstract base class defining core interface properties like MAC address, speed, MTU.
- Requires subclasses implement key methods like send/receive frames, enable/disable interface.
- Establishes universal network interface capabilities.
- Malicious Network Events Monitoring:
* Enhances network interfaces with the capability to monitor and capture Malicious Network Events (MNEs) based on predefined criteria such as specific keywords or traffic patterns.
* Integrates Number of Malicious Network Events (NMNE) detection functionalities, leveraging configurable settings like ``capture_nmne``, `nmne_capture_keywords``, and observation mechanisms such as ``NicObservation`` to classify and record network anomalies.
* Offers an additional layer of security and data analysis, crucial for identifying and mitigating malicious activities within the network infrastructure. Provides vital information for network security analysis and reinforcement learning algorithms.
**WiredNetworkInterface (Connection Type Layer)**

2
docs/source/simulation_components/network/nodes/firewall.rst
View File

@@ -229,7 +229,7 @@ To limit database server access to selected external IP addresses:
position=7
)
**Permitting DMZ Web Server Access while Blocking Specific Threats*
**Permitting DMZ Web Server Access while Blocking Specific Threats**
To authorize HTTP/HTTPS access to a DMZ-hosted web server, excluding known malicious IPs:

2
docs/source/simulation_components/network/nodes/network_node.rst
View File

@@ -27,7 +27,7 @@ in the transmission and routing of data within the simulated environment.
**Key Features:**
- **Frame Processing:** Central to the class is the ability to receive and process network frames, facilitating the
simulation of data flow through network devices.
simulation of data flow through network devices.
- **Abstract Methods:** Includes abstract methods such as ``receive_frame``, which subclasses must implement to specify
how devices handle incoming traffic.

105
docs/source/simulation_components/system/data_manipulation_bot.rst → docs/source/simulation_components/system/applications/data_manipulation_bot.rst
View File

@@ -2,20 +2,22 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DataManipulationBot:
DataManipulationBot
===================
###################
The ``DataManipulationBot`` class provides functionality to connect to a ``DatabaseService`` and execute malicious SQL statements.
The ``DataManipulationBot`` class provides functionality to connect to a :ref:`DatabaseService` and execute malicious SQL statements.
Overview
--------
========
The bot is intended to simulate a malicious actor carrying out attacks like:
- Dropping tables
- Deleting records
- Modifying data
on a database server by abusing an application's trusted database connectivity.
The bot performs attacks in the following stages to simulate the real pattern of an attack:
@@ -27,7 +29,7 @@ The bot performs attacks in the following stages to simulate the real pattern of
Each of these stages has a random, configurable probability of succeeding (by default 10%). The bot can also be configured to repeat the attack once complete.
Usage
-----
=====
- Create an instance and call ``configure`` to set:
- Target database server IP
@@ -40,34 +42,55 @@ The bot handles connecting, executing the statement, and disconnecting.
In a simulation, the bot can be controlled by using ``DataManipulationAgent`` which calls ``run`` on the bot at configured timesteps.
Example
-------
Implementation
==============
The bot connects to a :ref:`DatabaseClient` and leverages its connectivity. The host running ``DataManipulationBot`` must also have a :ref:`DatabaseClient` installed on it.
- Uses the Application base class for lifecycle management.
- Credentials, target IP and other options set via ``configure``.
- ``run`` handles connecting, executing statement, and disconnecting.
- SQL payload executed via ``query`` method.
- Results in malicious SQL being executed on remote database server.
Examples
========
Python
""""""
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.applications.database_client import DatabaseClient
client_1 = Computer(
hostname="client_1",
ip_address="192.168.10.21",
subnet_mask="255.255.255.0",
default_gateway="192.168.10.1"
default_gateway="192.168.10.1",
operating_state=NodeOperatingState.ON # initialise the computer in an ON state
)
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
client_1.software_manager.install(DataManipulationBot)
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
data_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DELETE")
data_manipulation_bot.run()
This would connect to the database service at 192.168.1.14, authenticate, and execute the SQL statement to drop the 'users' table.
This would connect to the database service at 192.168.1.14, authenticate, and execute the SQL statement to delete database contents.
Example with ``DataManipulationAgent``
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
""""""""""""""""""""""""""""""""""""""
If not using the data manipulation bot manually, it needs to be used with a data manipulation agent. Below is an example section of configuration file for setting up a simulation with data manipulation bot and agent.
.. code-block:: yaml
game_config:
game:
# ...
agents:
- ref: data_manipulation_red_bot
@@ -78,7 +101,7 @@ If not using the data manipulation bot manually, it needs to be used with a data
type: UC2RedObservation
options:
nodes:
- node_ref: client_1
- node_name: client_1
observations:
- logon_status
- operating_status
@@ -95,7 +118,7 @@ If not using the data manipulation bot manually, it needs to be used with a data
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_ref: client_1
- node_name: client_1
applications:
- application_ref: data_manipulation_bot
max_folders_per_node: 1
@@ -127,14 +150,56 @@ If not using the data manipulation bot manually, it needs to be used with a data
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
Implementation
--------------
Configuration
=============
The bot extends ``DatabaseClient`` and leverages its connectivity.
.. include:: ../common/common_configuration.rst
- Uses the Application base class for lifecycle management.
- Credentials, target IP and other options set via ``configure``.
- ``run`` handles connecting, executing statement, and disconnecting.
- SQL payload executed via ``query`` method.
- Results in malicious SQL being executed on remote database server.
.. |SOFTWARE_NAME| replace:: DataManipulationBot
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DataManipulationBot``
``server_ip``
"""""""""""""
IP address of the :ref:`DatabaseService` which the ``DataManipulationBot`` will try to attack.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``server_password``
"""""""""""""""""""
Optional. Default value is ``None``.
The password that the ``DataManipulationBot`` will use to access the :ref:`DatabaseService`.
``payload``
"""""""""""
Optional. Default value is ``DELETE``.
The payload that the ``DataManipulationBot`` will send to the :ref:`DatabaseService`.
.. include:: ../common/db_payload_list.rst
``port_scan_p_of_success``
""""""""""""""""""""""""""
Optional. Default value is ``0.1``.
The chance of the ``DataManipulationBot`` to succeed with a port scan (and therefore continue the attack).
This must be a float value between ``0`` and ``1``.
``data_manipulation_p_of_success``
""""""""""""""""""""""""""""""""""
Optional. Default value is ``0.1``.
The chance of the ``DataManipulationBot`` to succeed with a data manipulation attack.
This must be a float value between ``0`` and ``1``.

106
docs/source/simulation_components/system/applications/database_client.rst Normal file
View File

@@ -0,0 +1,106 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DatabaseClient:
DatabaseClient
##############
The ``DatabaseClient`` provides a client interface for connecting to the :ref:`DatabaseService`.
Key features
============
- Connects to the :ref:`DatabaseService` via the ``SoftwareManager``.
- Handles connecting and disconnecting.
- Executes SQL queries and retrieves result sets.
Usage
=====
- Initialise with server IP address and optional password.
- Connect to the :ref:`DatabaseService` with ``connect``.
- Retrieve results in a dictionary.
- Disconnect when finished.
Implementation
==============
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Connect and disconnect methods manage sessions.
- Payloads serialised as dictionaries for transmission.
- Extends base Application class.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.system.applications.database_client import DatabaseClient
client = Computer(
hostname="client",
ip_address="192.168.10.21",
subnet_mask="255.255.255.0",
default_gateway="192.168.10.1",
operating_state=NodeOperatingState.ON # initialise the computer in an ON state
)
# install DatabaseClient
client.software_manager.install(DatabaseClient)
database_client: DatabaseClient = client.software_manager.software.get("DatabaseClient")
# Configure the DatabaseClient
database_client.configure(server_ip_address=IPv4Address("192.168.0.1")) # address of the DatabaseService
database_client.run()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_computer
hostname: example_computer
type: computer
...
applications:
- ref: database_client
type: DatabaseClient
options:
db_server_ip: 192.168.0.1
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: DatabaseClient
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DatabaseClient``
``db_server_ip``
""""""""""""""""
IP address of the :ref:`DatabaseService` that the ``DatabaseClient`` will connect to
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``server_password``
"""""""""""""""""""
Optional. Default value is ``None``.
The password that the ``DatabaseClient`` will use to access the :ref:`DatabaseService`.

160
docs/source/simulation_components/system/applications/dos_bot.rst Normal file
View File

@@ -0,0 +1,160 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DoSBot:
DoSBot
######
The ``DoSBot`` is an implementation of a Denial of Service attack within the PrimAITE simulation. This specifically simulates a `Slow Loris attack <https://en.wikipedia.org/wiki/Slowloris_(computer_security)>`.
Key features
============
- Connects to the :ref:`DatabaseService` via the ``SoftwareManager``.
- Makes many connections to the :ref:`DatabaseService` which ends up using up the available connections.
Usage
=====
- Configure with target IP address and optional password.
- use ``run`` to run the application_loop of DoSBot to begin attacks
- DoSBot runs through different actions at each timestep
Implementation
==============
- Leverages :ref:`DatabaseClient` to create connections with :ref`DatabaseServer`.
- Extends base Application class.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
# Create Computer
computer = Computer(
hostname="computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
computer.power_on()
# Install DoSBot on computer
computer.software_manager.install(DoSBot)
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
# Configure the DoSBot
dos_bot.configure(
target_ip_address=IPv4Address("192.168.0.10"),
payload="SPOOF DATA",
repeat=False,
port_scan_p_of_success=0.8,
dos_intensity=1.0,
max_sessions=1000
)
# run DoSBot
dos_bot.run()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_computer
hostname: example_computer
type: computer
...
applications:
- ref: dos_bot
type: DoSBot
options:
target_ip_address: 192.168.0.10
payload: SPOOF DATA
repeat: False
port_scan_p_of_success: 0.8
dos_intensity: 1.0
max_sessions: 1000
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: DoSBot
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DoSBot``
``target_ip_address``
"""""""""""""""""""""
IP address of the :ref:`DatabaseService` which the ``DataManipulationBot`` will try to attack.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``target_port``
"""""""""""""""
Optional. Default value is ``5432``.
Port of the target service.
See :ref:`List of IPProtocols <List of IPProtocols>` for a list of protocols.
``payload``
"""""""""""
Optional. Default value is ``None``.
The payload that the ``DoSBot`` sends as part of its attack.
.. include:: ../common/db_payload_list.rst
``repeat``
""""""""""
Optional. Default value is ``False``.
If ``True`` the ``DoSBot`` will maintain its attack.
``port_scan_p_of_success``
""""""""""""""""""""""""""
Optional. Default value is ``0.1``.
The chance of the ``DoSBot`` to succeed with a port scan (and therefore continue the attack).
This must be a float value between ``0`` and ``1``.
``dos_intensity``
"""""""""""""""""
Optional. Default value is ``1.0``.
The intensity of the Denial of Service attack. This is multiplied by the number of ``max_sessions``.
This must be a float value between ``0`` and ``1``.
``max_sessions``
""""""""""""""""
Optional. Default value is ``1000``.
The maximum number of sessions the ``DoSBot`` is able to make.
This must be an integer value equal to or greater than ``0``.

111
docs/source/simulation_components/system/applications/web_browser.rst Normal file
View File

@@ -0,0 +1,111 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _WebBrowser:
WebBrowser
##########
The ``WebBrowser`` provides a client interface for connecting to the :ref:`WebServer`.
Key features
============
- Connects to the :ref:`WebServer` via the ``SoftwareManager``.
- Simulates HTTP requests and HTTP packet transfer across a network
- Allows the emulation of HTTP requests between client and server:
- Automatically uses ``DNSClient`` to resolve domain names
- GET: performs an HTTP GET request from client to server
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the ``WebBrowser``.
- Service runs on HTTP port 80 by default. (TODO: HTTPS)
- Execute sending an HTTP GET request with ``get_webpage``
Implementation
==============
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for making HTTP requests between an HTTP client and server.
- Extends base Service class.
Examples
========
Python
""""""
The ``WebBrowser`` utilises :ref:`DNSClient` and :ref:`DNSServer` to resolve a URL.
The :ref:`DNSClient` must be configured to use the :ref:`DNSServer`. The :ref:`DNSServer` should be configured to have the ``WebBrowser`` ``target_url`` within its ``domain_mapping``.
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.system.applications.web_browser import WebBrowser
# Create Computer
computer = Computer(
hostname="computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
computer.power_on()
# Install WebBrowser on computer
computer.software_manager.install(WebBrowser)
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
web_browser.run()
# configure the WebBrowser
web_browser.target_url = "arcd.com"
# once DNS server is configured with the correct domain mapping
# this should work
web_browser.get_webpage()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_computer
hostname: example_computer
type: computer
...
applications:
- ref: web_browser
type: WebBrowser
options:
target_url: http://arcd.com/
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: WebBrowser
.. |SOFTWARE_NAME_BACKTICK| replace:: ``WebBrowser``
``target_url``
""""""""""""""
The URL that the ``WebBrowser`` will request when ``get_webpage`` is called without parameters.
The URL can be in any format so long as the domain is within it e.g.
The domain ``arcd.com`` can be matched by
- http://arcd.com/
- http://arcd.com/users/
- arcd.com

18
docs/source/simulation_components/system/common/common_configuration.rst Normal file
View File

@@ -0,0 +1,18 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``ref``
=======
Human readable name used as reference for the |SOFTWARE_NAME_BACKTICK|. Not used in code.
``type``
========
The type of software that should be added. To add |SOFTWARE_NAME| this must be |SOFTWARE_NAME_BACKTICK|.
``options``
===========
The configuration options are the attributes that fall under the options for an application.

11
docs/source/simulation_components/system/common/db_payload_list.rst Normal file
View File

@@ -0,0 +1,11 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _Database Payload List:
Available Database Payloads:
- ``SELECT``
- ``INSERT``
- ``DELETE``

71
docs/source/simulation_components/system/database_client_server.rst
View File

@@ -1,71 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Database Client Server
======================
Database Service
----------------
The ``DatabaseService`` provides a SQL database server simulation by extending the base Service class.
Key capabilities
^^^^^^^^^^^^^^^^
- Creates a database file in the ``Node`` 's ``FileSystem`` upon creation.
- Handles connecting clients by maintaining a dictionary of connections mapped to session IDs.
- Authenticates connections using a configurable password.
- Simulates ``SELECT``, ``DELETE`` and ``INSERT`` SQL queries.
- Returns query results and status codes back to clients.
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Clients connect, execute queries, and disconnect.
- Service runs on TCP port 5432 by default.
Implementation
^^^^^^^^^^^^^^
- Creates the database file within the node's file system.
- Manages client connections in a dictionary by session ID.
- Processes SQL queries.
- Returns results and status codes in a standard dictionary format.
- Extends Service class for integration with ``SoftwareManager``.
Database Client
---------------
The DatabaseClient provides a client interface for connecting to the ``DatabaseService``.
Key features
^^^^^^^^^^^^
- Connects to the ``DatabaseService`` via the ``SoftwareManager``.
- Handles connecting and disconnecting.
- Executes SQL queries and retrieves result sets.
Usage
^^^^^
- Initialise with server IP address and optional password.
- Connect to the ``DatabaseService`` with ``connect``.
- Retrieve results in a dictionary.
- Disconnect when finished.
To create database backups:
- Configure the backup server on the ``DatabaseService`` by providing the Backup server ``IPv4Address`` with ``configure_backup``
- Create a backup using ``backup_database``. This fails if the backup server is not configured.
- Restore a backup using ``restore_backup``. By default, this uses the database created via ``backup_database``.
Implementation
^^^^^^^^^^^^^^
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Connect and disconnect methods manage sessions.
- Payloads serialised as dictionaries for transmission.
- Extends base Application class.

56
docs/source/simulation_components/system/dns_client_server.rst
View File

@@ -1,56 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
DNS Client Server
=================
DNS Server
----------
Also known as a DNS Resolver, the ``DNSServer`` provides a DNS Server simulation by extending the base Service class.
Key capabilities
^^^^^^^^^^^^^^^^
- Simulates DNS requests and DNSPacket transfer across a network
- Registers domain names and the IP Address linked to the domain name
- Returns the IP address for a given domain name within a DNS Packet that a DNS Client can read
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future)
Implementation
^^^^^^^^^^^^^^
- DNS request and responses use a ``DNSPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
DNS Client
----------
The DNSClient provides a client interface for connecting to the ``DNSServer``.
Key features
^^^^^^^^^^^^
- Connects to the ``DNSServer`` via the ``SoftwareManager``.
- Executes DNS lookup requests and keeps a cache of known domain name IP addresses.
- Handles connection to DNSServer and querying for domain name IP addresses.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future)
- Execute domain name checks with ``check_domain_exists``.
- ``DNSClient`` will automatically add the IP Address of the domain into its cache
Implementation
^^^^^^^^^^^^^^
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to find IP addresses via domain names.
- Extends base Service class.

135
docs/source/simulation_components/system/ftp_client_server.rst
View File

@@ -1,135 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
FTP Client Server
=================
FTP Server
----------
Provides a FTP Client-Server simulation by extending the base Service class.
Key capabilities
^^^^^^^^^^^^^^^^
- Simulates FTP requests and FTPPacket transfer across a network
- Allows the emulation of FTP commands between an FTP client and server:
- STOR: stores a file from client to server
- RETR: retrieves a file from the FTP server
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the FTP server service.
- Service runs on FTP (command) port 21 by default. (TODO: look at in depth implementation of FTP PORT command)
Implementation
^^^^^^^^^^^^^^
- FTP request and responses use a ``FTPPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
FTP Client
----------
The ``FTPClient`` provides a client interface for connecting to the ``FTPServer``.
Key features
^^^^^^^^^^^^
- Connects to the ``FTPServer`` via the ``SoftwareManager``.
- Simulates FTP requests and FTPPacket transfer across a network
- Allows the emulation of FTP commands between an FTP client and server:
- PORT: specifies the port that server should connect to on the client (currently only uses ``Port.FTP``)
- STOR: stores a file from client to server
- RETR: retrieves a file from the FTP server
- QUIT: disconnect from server
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the FTP client service.
- Service runs on FTP (command) port 21 by default. (TODO: look at in depth implementation of FTP PORT command)
- Execute sending a file to the FTP server with ``send_file``
- Execute retrieving a file from the FTP server with ``request_file``
Implementation
^^^^^^^^^^^^^^
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to transfer files between each other.
- Extends base Service class.
Example Usage
-------------
Dependencies
^^^^^^^^^^^^
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
Example peer to peer network
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
net = Network()
pc1 = Computer(
hostname="pc1",
ip_address="120.10.10.10",
subnet_mask="255.255.255.0",
operating_state=NodeOperatingState.ON # initialise the computer in an ON state
)
srv = Server(
hostname="srv",
ip_address="120.10.10.20",
subnet_mask="255.255.255.0",
operating_state=NodeOperatingState.ON # initialise the server in an ON state
)
net.connect(pc1.network_interface[1], srv.network_interface[1])
Install the FTP Server
^^^^^^^^^^^^^^^^^^^^^^
FTP Client should be pre installed on nodes
.. code-block:: python
srv.software_manager.install(FTPServer)
ftpserv: FTPServer = srv.software_manager.software['FTPServer']
Setting up the FTP Server
^^^^^^^^^^^^^^^^^^^^^^^^^
Set up the FTP Server with a file that the client will need to retrieve
.. code-block:: python
srv.file_system.create_file('my_file.png')
Check that file was retrieved
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
client.request_file(
src_folder_name='root',
src_file_name='my_file.png',
dest_folder_name='root',
dest_file_name='test.png',
dest_ip_address=IPv4Address("120.10.10.20")
)
print(client.get_file(folder_name="root", file_name="test.png"))

15
docs/source/simulation_components/system/list_of_applications.rst Normal file
View File

@@ -0,0 +1,15 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. toctree::
:maxdepth: 1
:glob:
applications/*
More info :py:mod:`primaite.game.game.APPLICATION_TYPES_MAPPING`
.. include:: list_of_system_applications.rst
.. |SOFTWARE_TYPE| replace:: application

15
docs/source/simulation_components/system/list_of_services.rst Normal file
View File

@@ -0,0 +1,15 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. toctree::
:maxdepth: 1
:glob:
services/*
More info :py:mod:`primaite.game.game.SERVICE_TYPES_MAPPING`
.. include:: list_of_system_services.rst
.. |SOFTWARE_TYPE| replace:: service

16
docs/source/simulation_components/system/list_of_system_applications.rst Normal file
View File

@@ -0,0 +1,16 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``system applications``
"""""""""""""""""""""""
Some applications are pre installed on nodes - this is similar to how some applications are included with the Operating System.
The application may not be configured as needed, in which case, see the relevant application page.
The list of applications that are considered system software are:
- ``WebBrowser``
More info :py:mod:`primaite.simulator.network.hardware.nodes.host.host_node.HostNode.SYSTEM_SOFTWARE`

18
docs/source/simulation_components/system/list_of_system_services.rst Normal file
View File

@@ -0,0 +1,18 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
``system services``
"""""""""""""""""""
Some services are pre installed on nodes - this is similar to how some services are included with the Operating System.
The service may not be configured as needed, in which case, see the relevant service page.
The list of services that are considered system software are:
- ``DNSClient``
- ``FTPClient``
- ``NTPClient``
More info :py:mod:`primaite.simulator.network.hardware.nodes.host.host_node.HostNode.SYSTEM_SOFTWARE`

54
docs/source/simulation_components/system/ntp_client_server.rst
View File

@@ -1,54 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
NTP Client Server
=================
NTP Server
----------
The ``NTPServer`` provides a NTP Server simulation by extending the base Service class.
NTP Client
----------
The ``NTPClient`` provides a NTP Client simulation by extending the base Service class.
Key capabilities
^^^^^^^^^^^^^^^^
- Simulates NTP requests and NTPPacket transfer across a network
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on UDP port 123 by default.
Implementation
^^^^^^^^^^^^^^
- NTP request and responses use a ``NTPPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
NTP Client
----------
The NTPClient provides a client interface for connecting to the ``NTPServer``.
Key features
^^^^^^^^^^^^
- Connects to the ``NTPServer`` via the ``SoftwareManager``.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on UDP port 123 by default.
Implementation
^^^^^^^^^^^^^^
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to find IP addresses via domain names.
- Extends base Service class.

116
docs/source/simulation_components/system/services/database_service.rst Normal file
View File

@@ -0,0 +1,116 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DatabaseService:
DatabaseService
###############
The ``DatabaseService`` provides a SQL database server simulation by extending the base Service class.
Key capabilities
================
- Creates a database file in the ``FileSystem`` of the ``Node`` (which the ``DatabaseService`` is installed on) upon creation.
- Handles connecting clients by maintaining a dictionary of connections mapped to session IDs.
- Authenticates connections using a configurable password.
- Simulates ``SELECT``, ``DELETE`` and ``INSERT`` SQL queries.
- Returns query results and status codes back to clients.
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Clients connect, execute queries, and disconnect.
- Service runs on TCP port 5432 by default.
**Supported queries:**
* ``SELECT``: As long as the database file is in a ``GOOD`` health state, the db service will respond with a 200 status code.
* ``DELETE``: This query represents an attack, it will cause the database file to enter a ``COMPROMISED`` state, and return a status code 200.
* ``INSERT``: If the database service is in a healthy state, this will return a 200 status, if it's not in a healthy state it will return 404.
* ``SELECT * FROM pg_stat_activity``: This query represents something an admin would send to check the status of the database. If the database service is in a healthy state, it returns a 200 status code, otherwise a 401 status code.
Implementation
==============
- Creates the database file within the node's file system.
- Manages client connections in a dictionary by session ID.
- Processes SQL queries.
- Returns results and status codes in a standard dictionary format.
- Extends Service class for integration with ``SoftwareManager``.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.database.database_service import DatabaseService
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install DatabaseService on server
server.software_manager.install(DatabaseService)
db_service: DatabaseService = server.software_manager.software.get("DatabaseService")
db_service.start()
# configure DatabaseService
db_service.configure_backup(IPv4Address("192.168.0.10"))
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: database_service
type: DatabaseService
options:
backup_server_ip: 192.168.0.10
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: DatabaseService
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DatabaseService``
``backup_server_ip``
""""""""""""""""""""
Optional. Default value is ``None``.
The IP Address of the backup server that the ``DatabaseService`` will use to create backups of the database.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
``password``
""""""""""""
Optional. Default value is ``None``.
The password that needs to be provided by connecting clients in order to create a successful connection.

99
docs/source/simulation_components/system/services/dns_client.rst Normal file
View File

@@ -0,0 +1,99 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DNSClient:
DNSClient
#########
The DNSClient provides a client interface for connecting to the :ref:`DNSServer`.
Key features
============
- Connects to the :ref:`DNSServer` via the ``SoftwareManager``.
- Executes DNS lookup requests and keeps a cache of known domain name IP addresses.
- Handles connection to DNSServer and querying for domain name IP addresses.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future)
- Execute domain name checks with ``check_domain_exists``.
- ``DNSClient`` will automatically add the IP Address of the domain into its cache
Implementation
==============
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to find IP addresses via domain names.
- Extends base Service class.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.dns.dns_client import DNSClient
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install DNSClient on server
server.software_manager.install(DNSClient)
dns_client: DNSClient = server.software_manager.software.get("DNSClient")
dns_client.start()
# configure DatabaseService
dns_client.dns_server = IPv4Address("192.168.0.10")
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: dns_client
type: DNSClient
options:
dns_server: 192.168.0.10
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: DNSClient
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DNSClient``
``dns_server``
""""""""""""""
Optional. Default value is ``None``.
The IP Address of the :ref:`DNSServer`.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.

98
docs/source/simulation_components/system/services/dns_server.rst Normal file
View File

@@ -0,0 +1,98 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _DNSServer:
DNSServer
#########
Also known as a DNS Resolver, the ``DNSServer`` provides a DNS Server simulation by extending the base Service class.
Key capabilities
================
- Simulates DNS requests and DNSPacket transfer across a network
- Registers domain names and the IP Address linked to the domain name
- Returns the IP address for a given domain name within a DNS Packet that a DNS Client can read
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future)
Implementation
==============
- DNS request and responses use a ``DNSPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.dns.dns_server import DNSServer
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install DNSServer on server
server.software_manager.install(DNSServer)
dns_server: DNSServer = server.software_manager.software.get("DNSServer")
dns_server.start()
# configure DatabaseService
dns_server.dns_register("arcd.com", IPv4Address("192.168.10.10"))
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: dns_server
type: DNSServer
options:
domain_mapping:
arcd.com: 192.168.0.10
another-example.com: 192.168.10.10
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: DNSServer
.. |SOFTWARE_NAME_BACKTICK| replace:: ``DNSServer``
domain_mapping
""""""""""""""
Domain mapping takes the domain and IP Addresses as a key-value pairs i.e.
If the domain is "arcd.com" and the IP Address attributed to the domain is 192.168.0.10, then the value should be ``arcd.com: 192.168.0.10``
The key must be a string and the IP Address must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.

91
docs/source/simulation_components/system/services/ftp_client.rst Normal file
View File

@@ -0,0 +1,91 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _FTPClient:
FTPClient
#########
The ``FTPClient`` provides a client interface for connecting to the :ref:`FTPServer`.
Key features
============
- Connects to the :ref:`FTPServer` via the ``SoftwareManager``.
- Simulates FTP requests and FTPPacket transfer across a network
- Allows the emulation of FTP commands between an FTP client and server:
- PORT: specifies the port that server should connect to on the client (currently only uses ``Port.FTP``)
- STOR: stores a file from client to server
- RETR: retrieves a file from the FTP server
- QUIT: disconnect from server
- Leverages the Service base class for install/uninstall, status tracking, etc.
- :ref:`FTPClient` and ``FTPServer`` utilise port 21 (FTP) throughout all file transfer / request
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the FTP client service.
- Service runs on FTP (command) port 21 by default
- Execute sending a file to the FTP server with ``send_file``
- Execute retrieving a file from the FTP server with ``request_file``
Implementation
==============
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to transfer files between each other.
- Extends base Service class.
Examples
========
Python
""""""
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.10",
start_up_duration=0,
)
server.power_on()
# Install FTPClient on server
server.software_manager.install(FTPClient)
ftp_client: FTPClient = server.software_manager.software.get("FTPClient")
ftp_client.start()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: ftp_client
type: FTPClient
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: FTPClient
.. |SOFTWARE_NAME_BACKTICK| replace:: ``FTPClient``
**FTPClient has no configuration options**

94
docs/source/simulation_components/system/services/ftp_server.rst Normal file
View File

@@ -0,0 +1,94 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _FTPServer:
FTPServer
#########
Provides a FTP Client-Server simulation by extending the base Service class.
Key capabilities
================
- Simulates FTP requests and FTPPacket transfer across a network
- Allows the emulation of FTP commands between an FTP client and server:
- STOR: stores a file from client to server
- RETR: retrieves a file from the FTP server
- Leverages the Service base class for install/uninstall, status tracking, etc.
- :ref:`FTPClient` and ``FTPServer`` utilise port 21 (FTP) throughout all file transfer / request
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the FTP server service.
- Service runs on FTP (command) port 21 by default
Implementation
==============
- FTP request and responses use a ``FTPPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
Examples
========
Python
""""""
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install FTPServer on server
server.software_manager.install(FTPServer)
ftp_server: FTPServer = server.software_manager.software.get("FTPServer")
ftp_server.start()
ftp_server.server_password = "test"
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: ftp_server
type: FTPServer
options:
server_password: test
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: FTPServer
.. |SOFTWARE_NAME_BACKTICK| replace:: ``FTPServer``
``server_password``
"""""""""""""""""""
Optional. Default value is ``None``.
The password that needs to be provided by a connecting :ref:`FTPClient` in order to create a successful connection.

95
docs/source/simulation_components/system/services/ntp_client.rst Normal file
View File

@@ -0,0 +1,95 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _NTPClient:
NTPClient
#########
The NTPClient provides a client interface for connecting to the ``NTPServer``.
Key features
============
- Connects to the ``NTPServer`` via the ``SoftwareManager``.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on UDP port 123 by default.
Implementation
==============
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for Nodes to find IP addresses via domain names.
- Extends base Service class.
Examples
========
Python
""""""
.. code-block:: python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install NTPClient on server
server.software_manager.install(NTPClient)
ntp_client: NTPClient = server.software_manager.software.get("NTPClient")
ntp_client.start()
ntp_client.configure(ntp_server_ip_address=IPv4Address("192.168.0.10"))
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: ntp_client
type: NTPClient
options:
ntp_server_ip: 192.168.0.10
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: NTPClient
.. |SOFTWARE_NAME_BACKTICK| replace:: ``NTPClient``
``ntp_server_ip``
"""""""""""""""""
Optional. Default value is ``None``.
The IP address of an NTP Server which provides a time that the ``NTPClient`` can synchronise to.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.

86
docs/source/simulation_components/system/services/ntp_server.rst Normal file
View File

@@ -0,0 +1,86 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _NTPServer:
NTPServer
#########
The ``NTPServer`` provides a NTP Server simulation by extending the base Service class.
NTP Client
==========
The ``NTPClient`` provides a NTP Client simulation by extending the base Service class.
Key capabilities
================
- Simulates NTP requests and NTPPacket transfer across a network
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the database service.
- Service runs on UDP port 123 by default.
Implementation
==============
- NTP request and responses use a ``NTPPacket`` object
- Extends Service class for integration with ``SoftwareManager``.
Examples
========
Python
""""""
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install NTPServer on server
server.software_manager.install(NTPServer)
ntp_server: NTPServer = server.software_manager.software.get("NTPServer")
ntp_server.start()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: ntp_server
type: NTPServer
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: NTPServer
.. |SOFTWARE_NAME_BACKTICK| replace:: ``NTPServer``
**NTPServer has no configuration options**

86
docs/source/simulation_components/system/services/web_server.rst Normal file
View File

@@ -0,0 +1,86 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _WebServer:
WebServer
#########
Provides a Web Server simulation by extending the base Service class.
Key capabilities
================
- Simulates a web server with the capability to also request data from a database
- Allows the emulation of HTTP requests between client (e.g. a web browser) and server
- GET request sends a get all users request to the database server and returns an HTTP 200 status if the database is responsive
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
=====
- Install on a Node via the ``SoftwareManager`` to start the `WebServer`.
- Service runs on HTTP port 80 by default. (TODO: HTTPS)
- A :ref:`DatabaseClient` must be installed and configured on the same node as the ``WebServer`` if it is intended to send a users request i.e.
in the case that the :ref:`WebBrowser` sends a request with users in its request path, the ``WebServer`` will utilise the ``DatabaseClient`` to send a request to the ``DatabaseService``
Implementation
==============
- HTTP request uses a ``HttpRequestPacket`` object
- HTTP response uses a ``HttpResponsePacket`` object
- Extends Service class for integration with ``SoftwareManager``.
Examples
========
Python
""""""
.. code-block:: python
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.web_server.web_server import WebServer
# Create Server
server = Server(
hostname="server",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server.power_on()
# Install WebServer on server
server.software_manager.install(WebServer)
web_server: WebServer = server.software_manager.software.get("WebServer")
web_server.start()
Via Configuration
"""""""""""""""""
.. code-block:: yaml
simulation:
network:
nodes:
- ref: example_server
hostname: example_server
type: server
...
services:
- ref: web_server
type: WebServer
Configuration
=============
.. include:: ../common/common_configuration.rst
.. |SOFTWARE_NAME| replace:: WebServer
.. |SOFTWARE_NAME_BACKTICK| replace:: ``WebServer``
**WebServer has no configuration options**

2
docs/source/simulation_components/system/session_and_software_manager.rst
View File

@@ -16,6 +16,8 @@ ARP, ICMP, or the Web Client. This pathway exemplifies the structured processing
each frame reaches its intended target within the simulated environment.
.. image:: node_session_software_model_example.png
:width: 500
:align: center
Session Manager
---------------

32
docs/source/simulation_components/system/software.rst
View File

@@ -10,7 +10,7 @@ Software
Base Software
-------------
All software which inherits ``IOSoftware`` installed on a node will not work unless the node has been turned on.
Software which inherits ``IOSoftware`` installed on a node will not work unless the node has been turned on.
See :ref:`Node Start up and Shut down`
@@ -39,15 +39,27 @@ See :ref:`Node Start up and Shut down`
assert node.operating_state is NodeOperatingState.ON
assert web_server.operating_state is ServiceOperatingState.RUNNING # service turned back on when node is powered on
.. _List of Applications:
Services, Processes and Applications:
#####################################
Applications
############
.. toctree::
:maxdepth: 2
These are a list of applications that are currently available in PrimAITE:
database_client_server
data_manipulation_bot
dns_client_server
ftp_client_server
web_browser_and_web_server_service
.. include:: list_of_applications.rst
.. _List of Services:
Services
########
These are a list of services that are currently available in PrimAITE:
.. include:: list_of_services.rst
.. _List of Processes:
Processes
#########
`To be implemented`

110
docs/source/simulation_components/system/web_browser_and_web_server_service.rst
View File

@@ -1,110 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Web Browser and Web Server Service
==================================
Web Server Service
------------------
Provides a Web Server simulation by extending the base Service class.
Key capabilities
^^^^^^^^^^^^^^^^
- Simulates a web server with the capability to also request data from a database
- Allows the emulation of HTTP requests between client (e.g. a web browser) and server
- GET request sends a get all users request to the database server and returns an HTTP 200 status if the database is responsive
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the `WebServer`.
- Service runs on HTTP port 80 by default. (TODO: HTTPS)
Implementation
^^^^^^^^^^^^^^
- HTTP request uses a ``HttpRequestPacket`` object
- HTTP response uses a ``HttpResponsePacket`` object
- Extends Service class for integration with ``SoftwareManager``.
Web Browser (Web Client)
------------------------
The ``WebBrowser`` provides a client interface for connecting to the ``WebServer``.
Key features
^^^^^^^^^^^^
- Connects to the ``WebServer`` via the ``SoftwareManager``.
- Simulates HTTP requests and HTTP packet transfer across a network
- Allows the emulation of HTTP requests between client and server:
- Automatically uses ``DNSClient`` to resolve domain names
- GET: performs an HTTP GET request from client to server
- Leverages the Service base class for install/uninstall, status tracking, etc.
Usage
^^^^^
- Install on a Node via the ``SoftwareManager`` to start the ``WebBrowser``.
- Service runs on HTTP port 80 by default. (TODO: HTTPS)
- Execute sending an HTTP GET request with ``get_webpage``
Implementation
^^^^^^^^^^^^^^
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Provides easy interface for making HTTP requests between an HTTP client and server.
- Extends base Service class.
Example Usage
-------------
Dependencies
^^^^^^^^^^^^
.. code-block:: python
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.web_server.web_server_service import WebServer
Example peer to peer network
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
net = Network()
pc1 = Computer(hostname="pc1", ip_address="192.168.1.50", subnet_mask="255.255.255.0")
srv = Server(hostname="srv", ip_address="192.168.1.10", subnet_mask="255.255.255.0")
pc1.power_on()
srv.power_on()
net.connect(pc1.network_interface[1], srv.network_interface[1])
Install the Web Server
^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
# web browser is automatically installed in computer nodes
# IRL this is usually included with an OS
client: WebBrowser = pc1.software_manager.software['WebBrowser']
# install web server
srv.software_manager.install(WebServer)
webserv: WebServer = srv.software_manager.software['WebServer']
Open the web page
^^^^^^^^^^^^^^^^^
Using a domain name to connect to a website requires setting up DNS Servers. For this example, it is possible to use the IP address directly
.. code-block:: python
# check that the get request succeeded
print(client.get_webpage("http://192.168.1.10")) # should be True

11
docs/source/simulation_structure.rst
View File

@@ -12,14 +12,15 @@ and a domain controller for managing software and users.
Each node of the simulation 'tree' has responsibility for creating, deleting, and updating its direct descendants. Also,
when a component's ``describe_state()`` method is called, it will include the state of its descendants. The
``apply_request()`` method can be used to act on a component or one of its descendatnts. The diagram below shows the
``apply_request()`` method can be used to act on a component or one of its descendants. The diagram below shows the
relationship between components.
.. image:: _static/component_relationship.png
.. image:: ../../_static/component_relationship.png
:width: 500
:alt: The top level simulation object owns a NetworkContainer and a DomainController. The DomainController has a
list of accounts. The network container has links and nodes. Nodes can own switchports, NICs, FileSystem,
Application, Service, and Process.
:align: center
:alt: :: The top level simulation object owns a NetworkContainer and a DomainController. The DomainController has a
list of accounts. The network container has links and nodes. Nodes can own switchports, NICs, FileSystem,
Application, Service, and Process.
Actions

6
docs/source/state_system.rst
View File

@@ -3,11 +3,11 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Simulation State
==============
================
``SimComponent`` objects in the simulation have a method called ``describe_state`` which return a dictionary of the state of the component. This is used to report pertinent data that could impact an agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` object reports not only its own attributes in the state but also those of its child components. I.e. a computer node will report the state of its ``FileSystem`` and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling the childrens' own ``describe_state`` methods.
``SimComponent`` objects in the simulation have a method called ``describe_state`` which return a dictionary of the state of the component. This is used to report pertinent data that could impact an agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` object reports not only its own attributes in the state but also those of its child components. I.e. a computer node will report the state of its ``FileSystem`` and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling the children's own ``describe_state`` methods.
The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then passes the state to the agents once per simulation step. For this reason, all ``SimComponent`` objetcs must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``.
The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then passes the state to the agents once per simulation step. For this reason, all ``SimComponent`` objects must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``.
This code snippet demonstrates how the state information is defined within the ``SimComponent`` class:

4
pyproject.toml
View File

@@ -57,7 +57,7 @@ dev = [
"build==0.10.0",
"flake8==6.0.0",
"flake8-annotations",
"furo==2023.3.27",
"furo==2024.01.29",
"gputil==1.4.0",
"pip-licenses==4.3.0",
"pre-commit==2.20.0",
@@ -67,7 +67,7 @@ dev = [
"pytest-cov==4.0.0",
"pytest-flake8==1.1.1",
"setuptools==66",
"Sphinx==6.1.3",
"Sphinx==7.1.2",
"sphinx-copybutton==0.5.2",
"wheel==0.38.4"
]

4
src/primaite/cli.py
View File

@@ -127,10 +127,10 @@ def session(
:param config: The path to the config file. Optional, if None, the example config will be used.
:type config: Optional[str]
"""
from primaite.config.load import example_config_path
from primaite.config.load import data_manipulation_config_path
from primaite.main import run
if not config:
config = example_config_path()
config = data_manipulation_config_path()
print(config)
run(config_path=config, agent_load_path=agent_load_file)

960
src/primaite/config/_package_data/data_manipulation.yaml Normal file
View File

@@ -0,0 +1,960 @@
training_config:
rl_framework: SB3
rl_algorithm: PPO
seed: 333
n_learn_episodes: 1
n_eval_episodes: 5
max_steps_per_episode: 128
deterministic_eval: false
n_agents: 1
agent_references:
- defender
io_settings:
save_agent_actions: true
save_step_metadata: false
save_pcap_logs: false
save_sys_logs: false
game:
max_episode_length: 128
ports:
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
thresholds:
nmne:
high: 10
medium: 5
low: 0
agents:
- ref: client_2_green_user
team: GREEN
type: ProbabilisticAgent
agent_settings:
action_probabilities:
0: 0.3
1: 0.6
2: 0.1
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_2
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
options: {}
1:
action: NODE_APPLICATION_EXECUTE
options:
node_id: 0
application_id: 0
2:
action: NODE_APPLICATION_EXECUTE
options:
node_id: 0
application_id: 1
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.25
options:
node_hostname: client_2
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
weight: 0.05
options:
node_hostname: client_2
- ref: client_1_green_user
team: GREEN
type: ProbabilisticAgent
agent_settings:
action_probabilities:
0: 0.3
1: 0.6
2: 0.1
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
options: {}
1:
action: NODE_APPLICATION_EXECUTE
options:
node_id: 0
application_id: 0
2:
action: NODE_APPLICATION_EXECUTE
options:
node_id: 0
application_id: 1
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.25
options:
node_hostname: client_1
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
weight: 0.05
options:
node_hostname: client_1
- ref: data_manipulation_attacker
team: RED
type: RedDatabaseCorruptingAgent
observation_space:
type: UC2RedObservation
options:
nodes: {}
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: DataManipulationBot
- node_name: client_2
applications:
- application_name: DataManipulationBot
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
reward_function:
reward_components:
- type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self)
start_settings:
start_step: 25
frequency: 20
variance: 5
- ref: defender
team: BLUE
type: ProxyAgent
observation_space:
type: UC2BlueObservation
options:
num_services_per_node: 1
num_folders_per_node: 1
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_hostname: domain_controller
services:
- service_name: DNSServer
- node_hostname: web_server
services:
- service_name: WebServer
- node_hostname: database_server
folders:
- folder_name: database
files:
- file_name: database.db
- node_hostname: backup_server
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
- link_ref: switch_1___domain_controller
- link_ref: switch_1___web_server
- link_ref: switch_1___database_server
- link_ref: switch_1___backup_server
- link_ref: switch_1___security_suite
- link_ref: switch_2___client_1
- link_ref: switch_2___client_2
- link_ref: switch_2___security_suite
acl:
options:
max_acl_rules: 10
router_hostname: router_1
ip_address_order:
- node_hostname: domain_controller
nic_num: 1
- node_hostname: web_server
nic_num: 1
- node_hostname: database_server
nic_num: 1
- node_hostname: backup_server
nic_num: 1
- node_hostname: security_suite
nic_num: 1
- node_hostname: client_1
nic_num: 1
- node_hostname: client_2
nic_num: 1
- node_hostname: security_suite
nic_num: 2
ics: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_SERVICE_SCAN
- type: NODE_SERVICE_STOP
- type: NODE_SERVICE_START
- type: NODE_SERVICE_PAUSE
- type: NODE_SERVICE_RESUME
- type: NODE_SERVICE_RESTART
- type: NODE_SERVICE_DISABLE
- type: NODE_SERVICE_ENABLE
- type: NODE_SERVICE_PATCH
- type: NODE_FILE_SCAN
- type: NODE_FILE_CHECKHASH
- type: NODE_FILE_DELETE
- type: NODE_FILE_REPAIR
- type: NODE_FILE_RESTORE
- type: NODE_FOLDER_SCAN
- type: NODE_FOLDER_CHECKHASH
- type: NODE_FOLDER_REPAIR
- type: NODE_FOLDER_RESTORE
- type: NODE_OS_SCAN
- type: NODE_SHUTDOWN
- type: NODE_STARTUP
- type: NODE_RESET
- type: NETWORK_ACL_ADDRULE
options:
target_router_hostname: router_1
- type: NETWORK_ACL_REMOVERULE
options:
target_router_hostname: router_1
- type: NETWORK_NIC_ENABLE
- type: NETWORK_NIC_DISABLE
action_map:
0:
action: DONOTHING
options: {}
# scan webapp service
1:
action: NODE_SERVICE_SCAN
options:
node_id: 1
service_id: 0
# stop webapp service
2:
action: NODE_SERVICE_STOP
options:
node_id: 1
service_id: 0
# start webapp service
3:
action: "NODE_SERVICE_START"
options:
node_id: 1
service_id: 0
4:
action: "NODE_SERVICE_PAUSE"
options:
node_id: 1
service_id: 0
5:
action: "NODE_SERVICE_RESUME"
options:
node_id: 1
service_id: 0
6:
action: "NODE_SERVICE_RESTART"
options:
node_id: 1
service_id: 0
7:
action: "NODE_SERVICE_DISABLE"
options:
node_id: 1
service_id: 0
8:
action: "NODE_SERVICE_ENABLE"
options:
node_id: 1
service_id: 0
9: # check database.db file
action: "NODE_FILE_SCAN"
options:
node_id: 2
folder_id: 0
file_id: 0
10:
action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
options:
node_id: 2
folder_id: 0
file_id: 0
11:
action: "NODE_FILE_DELETE"
options:
node_id: 2
folder_id: 0
file_id: 0
12:
action: "NODE_FILE_REPAIR"
options:
node_id: 2
folder_id: 0
file_id: 0
13:
action: "NODE_SERVICE_PATCH"
options:
node_id: 2
service_id: 0
14:
action: "NODE_FOLDER_SCAN"
options:
node_id: 2
folder_id: 0
15:
action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
options:
node_id: 2
folder_id: 0
16:
action: "NODE_FOLDER_REPAIR"
options:
node_id: 2
folder_id: 0
17:
action: "NODE_FOLDER_RESTORE"
options:
node_id: 2
folder_id: 0
18:
action: "NODE_OS_SCAN"
options:
node_id: 0
19:
action: "NODE_SHUTDOWN"
options:
node_id: 0
20:
action: NODE_STARTUP
options:
node_id: 0
21:
action: NODE_RESET
options:
node_id: 0
22:
action: "NODE_OS_SCAN"
options:
node_id: 1
23:
action: "NODE_SHUTDOWN"
options:
node_id: 1
24:
action: NODE_STARTUP
options:
node_id: 1
25:
action: NODE_RESET
options:
node_id: 1
26: # old action num: 18
action: "NODE_OS_SCAN"
options:
node_id: 2
27:
action: "NODE_SHUTDOWN"
options:
node_id: 2
28:
action: NODE_STARTUP
options:
node_id: 2
29:
action: NODE_RESET
options:
node_id: 2
30:
action: "NODE_OS_SCAN"
options:
node_id: 3
31:
action: "NODE_SHUTDOWN"
options:
node_id: 3
32:
action: NODE_STARTUP
options:
node_id: 3
33:
action: NODE_RESET
options:
node_id: 3
34:
action: "NODE_OS_SCAN"
options:
node_id: 4
35:
action: "NODE_SHUTDOWN"
options:
node_id: 4
36:
action: NODE_STARTUP
options:
node_id: 4
37:
action: NODE_RESET
options:
node_id: 4
38:
action: "NODE_OS_SCAN"
options:
node_id: 5
39: # old action num: 19 # shutdown client 1
action: "NODE_SHUTDOWN"
options:
node_id: 5
40: # old action num: 20
action: NODE_STARTUP
options:
node_id: 5
41: # old action num: 21
action: NODE_RESET
options:
node_id: 5
42:
action: "NODE_OS_SCAN"
options:
node_id: 6
43:
action: "NODE_SHUTDOWN"
options:
node_id: 6
44:
action: NODE_STARTUP
options:
node_id: 6
45:
action: NODE_RESET
options:
node_id: 6
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
permission: 2
source_ip_id: 7 # client 1
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "NETWORK_ACL_ADDRULE"
options:
position: 2
permission: 2
source_ip_id: 8 # client 2
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 3
permission: 2
source_ip_id: 7 # client 1
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 4
permission: 2
source_ip_id: 8 # client 2
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
50: # old action num: 26
action: "NETWORK_ACL_ADDRULE"
options:
position: 5
permission: 2
source_ip_id: 7 # client 1
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
51: # old action num: 27
action: "NETWORK_ACL_ADDRULE"
options:
position: 6
permission: 2
source_ip_id: 8 # client 2
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
52: # old action num: 28
action: "NETWORK_ACL_REMOVERULE"
options:
position: 0
53: # old action num: 29
action: "NETWORK_ACL_REMOVERULE"
options:
position: 1
54: # old action num: 30
action: "NETWORK_ACL_REMOVERULE"
options:
position: 2
55: # old action num: 31
action: "NETWORK_ACL_REMOVERULE"
options:
position: 3
56: # old action num: 32
action: "NETWORK_ACL_REMOVERULE"
options:
position: 4
57: # old action num: 33
action: "NETWORK_ACL_REMOVERULE"
options:
position: 5
58: # old action num: 34
action: "NETWORK_ACL_REMOVERULE"
options:
position: 6
59: # old action num: 35
action: "NETWORK_ACL_REMOVERULE"
options:
position: 7
60: # old action num: 36
action: "NETWORK_ACL_REMOVERULE"
options:
position: 8
61: # old action num: 37
action: "NETWORK_ACL_REMOVERULE"
options:
position: 9
62: # old action num: 38
action: "NETWORK_NIC_DISABLE"
options:
node_id: 0
nic_id: 0
63: # old action num: 39
action: "NETWORK_NIC_ENABLE"
options:
node_id: 0
nic_id: 0
64: # old action num: 40
action: "NETWORK_NIC_DISABLE"
options:
node_id: 1
nic_id: 0
65: # old action num: 41
action: "NETWORK_NIC_ENABLE"
options:
node_id: 1
nic_id: 0
66: # old action num: 42
action: "NETWORK_NIC_DISABLE"
options:
node_id: 2
nic_id: 0
67: # old action num: 43
action: "NETWORK_NIC_ENABLE"
options:
node_id: 2
nic_id: 0
68: # old action num: 44
action: "NETWORK_NIC_DISABLE"
options:
node_id: 3
nic_id: 0
69: # old action num: 45
action: "NETWORK_NIC_ENABLE"
options:
node_id: 3
nic_id: 0
70: # old action num: 46
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 0
71: # old action num: 47
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 0
72: # old action num: 48
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 1
73: # old action num: 49
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 1
74: # old action num: 50
action: "NETWORK_NIC_DISABLE"
options:
node_id: 5
nic_id: 0
75: # old action num: 51
action: "NETWORK_NIC_ENABLE"
options:
node_id: 5
nic_id: 0
76: # old action num: 52
action: "NETWORK_NIC_DISABLE"
options:
node_id: 6
nic_id: 0
77: # old action num: 53
action: "NETWORK_NIC_ENABLE"
options:
node_id: 6
nic_id: 0
options:
nodes:
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_address_order:
- node_name: domain_controller
nic_num: 1
- node_name: web_server
nic_num: 1
- node_name: database_server
nic_num: 1
- node_name: backup_server
nic_num: 1
- node_name: security_suite
nic_num: 1
- node_name: client_1
nic_num: 1
- node_name: client_2
nic_num: 1
- node_name: security_suite
nic_num: 2
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
weight: 0.40
options:
node_hostname: database_server
folder_name: database
file_name: database.db
- type: SHARED_REWARD
weight: 1.0
options:
agent_name: client_1_green_user
- type: SHARED_REWARD
weight: 1.0
options:
agent_name: client_2_green_user
agent_settings:
flatten_obs: true
simulation:
network:
nmne_config:
capture_nmne: true
nmne_capture_keywords:
- DELETE
nodes:
- ref: router_1
hostname: router_1
type: router
num_ports: 5
ports:
1:
ip_address: 192.168.1.1
subnet_mask: 255.255.255.0
2:
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
acl:
18:
action: PERMIT
src_port: POSTGRES_SERVER
dst_port: POSTGRES_SERVER
19:
action: PERMIT
src_port: DNS
dst_port: DNS
20:
action: PERMIT
src_port: FTP
dst_port: FTP
21:
action: PERMIT
src_port: HTTP
dst_port: HTTP
22:
action: PERMIT
src_port: ARP
dst_port: ARP
23:
action: PERMIT
protocol: ICMP
- ref: switch_1
hostname: switch_1
type: switch
num_ports: 8
- ref: switch_2
hostname: switch_2
type: switch
num_ports: 8
- ref: domain_controller
hostname: domain_controller
type: server
ip_address: 192.168.1.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
services:
- ref: domain_controller_dns_server
type: DNSServer
options:
domain_mapping:
arcd.com: 192.168.1.12 # web server
- ref: web_server
hostname: web_server
type: server
ip_address: 192.168.1.12
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: database_server
hostname: database_server
type: server
ip_address: 192.168.1.14
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: database_service
type: DatabaseService
options:
backup_server_ip: 192.168.1.16
- ref: database_ftp_client
type: FTPClient
- ref: backup_server
hostname: backup_server
type: server
ip_address: 192.168.1.16
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: backup_service
type: FTPServer
- ref: security_suite
hostname: security_suite
type: server
ip_address: 192.168.1.110
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
network_interfaces:
2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot
ip_address: 192.168.10.110
subnet_mask: 255.255.255.0
- ref: client_1
hostname: client_1
type: computer
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- ref: data_manipulation_bot
type: DataManipulationBot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- ref: client_1_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
- ref: client_1_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
services:
- ref: client_1_dns_client
type: DNSClient
- ref: client_2
hostname: client_2
type: computer
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- ref: client_2_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
- ref: data_manipulation_bot
type: DataManipulationBot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- ref: client_2_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
services:
- ref: client_2_dns_client
type: DNSClient
links:
- ref: router_1___switch_1
endpoint_a_ref: router_1
endpoint_a_port: 1
endpoint_b_ref: switch_1
endpoint_b_port: 8
- ref: router_1___switch_2
endpoint_a_ref: router_1
endpoint_a_port: 2
endpoint_b_ref: switch_2
endpoint_b_port: 8
- ref: switch_1___domain_controller
endpoint_a_ref: switch_1
endpoint_a_port: 1
endpoint_b_ref: domain_controller
endpoint_b_port: 1
- ref: switch_1___web_server
endpoint_a_ref: switch_1
endpoint_a_port: 2
endpoint_b_ref: web_server
endpoint_b_port: 1
- ref: switch_1___database_server
endpoint_a_ref: switch_1
endpoint_a_port: 3
endpoint_b_ref: database_server
endpoint_b_port: 1
- ref: switch_1___backup_server
endpoint_a_ref: switch_1
endpoint_a_port: 4
endpoint_b_ref: backup_server
endpoint_b_port: 1
- ref: switch_1___security_suite
endpoint_a_ref: switch_1
endpoint_a_port: 7
endpoint_b_ref: security_suite
endpoint_b_port: 1
- ref: switch_2___client_1
endpoint_a_ref: switch_2
endpoint_a_port: 1
endpoint_b_ref: client_1
endpoint_b_port: 1
- ref: switch_2___client_2
endpoint_a_ref: switch_2
endpoint_a_port: 2
endpoint_b_ref: client_2
endpoint_b_port: 1
- ref: switch_2___security_suite
endpoint_a_ref: switch_2
endpoint_a_port: 7
endpoint_b_ref: security_suite
endpoint_b_port: 2

1037
src/primaite/config/_package_data/example_config_2_rl_agents.yaml → src/primaite/config/_package_data/data_manipulation_marl.yaml
View File

File diff suppressed because it is too large Load Diff

4
src/primaite/config/load.py
View File

@@ -30,14 +30,14 @@ def load(file_path: Union[str, Path]) -> Dict:
return config
def example_config_path() -> Path:
def data_manipulation_config_path() -> Path:
"""
Get the path to the example config.
:return: Path to the example config.
:rtype: Path
"""
path = _EXAMPLE_CFG / "example_config.yaml"
path = _EXAMPLE_CFG / "data_manipulation.yaml"
if not path.exists():
msg = f"Example config does not exist: {path}. Have you run `primaite setup`?"
_LOGGER.error(msg)

111
src/primaite/game/agent/actions.py
View File

@@ -492,9 +492,9 @@ class NetworkACLAddRuleAction(AbstractAction):
"add_rule",
permission_str,
protocol,
src_ip,
str(src_ip),
src_port,
dst_ip,
str(dst_ip),
dst_port,
position,
]
@@ -572,7 +572,7 @@ class NetworkNICDisableAction(NetworkNICAbstractAction):
class ActionManager:
"""Class which manages the action space for an agent."""
_act_class_identifiers: Dict[str, type] = {
act_class_identifiers: Dict[str, type] = {
"DONOTHING": DoNothingAction,
"NODE_SERVICE_SCAN": NodeServiceScanAction,
"NODE_SERVICE_STOP": NodeServiceStopAction,
@@ -607,7 +607,6 @@ class ActionManager:
def __init__(
self,
game: "PrimaiteGame", # reference to game for information lookup
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
max_folders_per_node: int = 2, # allows calculating shape
@@ -618,7 +617,7 @@ class ActionManager:
max_acl_rules: int = 10, # allows calculating shape
protocols: List[str] = ["TCP", "UDP", "ICMP"], # allow mapping index to protocol
ports: List[str] = ["HTTP", "DNS", "ARP", "FTP", "NTP"], # allow mapping index to port
ip_address_list: Optional[List[str]] = None, # to allow us to map an index to an ip address.
ip_address_list: List[str] = [], # to allow us to map an index to an ip address.
act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions
) -> None:
"""Init method for ActionManager.
@@ -649,7 +648,6 @@ class ActionManager:
:param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
:type act_map: Optional[Dict[int, Dict]]
"""
self.game: "PrimaiteGame" = game
self.node_names: List[str] = [n["node_name"] for n in nodes]
"""List of node names in this action space. The list order is the mapping between node index and node name."""
self.application_names: List[List[str]] = []
@@ -707,25 +705,7 @@ class ActionManager:
self.protocols: List[str] = protocols
self.ports: List[str] = ports
self.ip_address_list: List[str]
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
if ip_address_list is not None:
self.ip_address_list = ip_address_list
else:
self.ip_address_list = []
for node_name in self.node_names:
node_obj = self.game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
self.ip_address_list.append(nic_obj.ip_address)
self.ip_address_list: List[str] = ip_address_list
# action_args are settings which are applied to the action space as a whole.
global_action_args = {
@@ -753,7 +733,7 @@ class ActionManager:
# and `options` is an optional dict of options to pass to the init method of the action class
act_type = act_spec.get("type")
act_options = act_spec.get("options", {})
self.actions[act_type] = self._act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.actions[act_type] = self.act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.action_map: Dict[int, Tuple[str, Dict]] = {}
"""
@@ -832,6 +812,13 @@ class ActionManager:
:return: The node hostname.
:rtype: str
"""
if not node_idx < len(self.node_names):
msg = (
f"Error: agent attempted to perform an action on node {node_idx}, but its action space only"
f"has {len(self.node_names)} nodes."
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.node_names[node_idx]
def get_folder_name_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]:
@@ -845,6 +832,13 @@ class ActionManager:
:return: The name of the folder. Or None if the node has fewer folders than the given index.
:rtype: Optional[str]
"""
if node_idx >= len(self.folder_names) or folder_idx >= len(self.folder_names[node_idx]):
msg = (
f"Error: agent attempted to perform an action on node {node_idx} and folder {folder_idx}, but this"
f" is out of range for its action space. Folder on each node: {self.folder_names}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.folder_names[node_idx][folder_idx]
def get_file_name_by_idx(self, node_idx: int, folder_idx: int, file_idx: int) -> Optional[str]:
@@ -860,6 +854,17 @@ class ActionManager:
fewer files than the given index.
:rtype: Optional[str]
"""
if (
node_idx >= len(self.file_names)
or folder_idx >= len(self.file_names[node_idx])
or file_idx >= len(self.file_names[node_idx][folder_idx])
):
msg = (
f"Error: agent attempted to perform an action on node {node_idx} folder {folder_idx} file {file_idx}"
f" but this is out of range for its action space. Files on each node: {self.file_names}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.file_names[node_idx][folder_idx][file_idx]
def get_service_name_by_idx(self, node_idx: int, service_idx: int) -> Optional[str]:
@@ -872,6 +877,13 @@ class ActionManager:
:return: The name of the service. Or None if the node has fewer services than the given index.
:rtype: Optional[str]
"""
if node_idx >= len(self.service_names) or service_idx >= len(self.service_names[node_idx]):
msg = (
f"Error: agent attempted to perform an action on node {node_idx} and service {service_idx}, but this"
f" is out of range for its action space. Services on each node: {self.service_names}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.service_names[node_idx][service_idx]
def get_application_name_by_idx(self, node_idx: int, application_idx: int) -> Optional[str]:
@@ -884,6 +896,13 @@ class ActionManager:
:return: The name of the service. Or None if the node has fewer services than the given index.
:rtype: Optional[str]
"""
if node_idx >= len(self.application_names) or application_idx >= len(self.application_names[node_idx]):
msg = (
f"Error: agent attempted to perform an action on node {node_idx} and app {application_idx}, but "
f"this is out of range for its action space. Applications on each node: {self.application_names}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.application_names[node_idx][application_idx]
def get_internet_protocol_by_idx(self, protocol_idx: int) -> str:
@@ -894,6 +913,13 @@ class ActionManager:
:return: The protocol.
:rtype: str
"""
if protocol_idx >= len(self.protocols):
msg = (
f"Error: agent attempted to perform an action on protocol {protocol_idx} but this"
f" is out of range for its action space. Protocols: {self.protocols}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.protocols[protocol_idx]
def get_ip_address_by_idx(self, ip_idx: int) -> str:
@@ -905,6 +931,13 @@ class ActionManager:
:return: The IP address.
:rtype: str
"""
if ip_idx >= len(self.ip_address_list):
msg = (
f"Error: agent attempted to perform an action on ip address {ip_idx} but this"
f" is out of range for its action space. IP address list: {self.ip_address_list}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.ip_address_list[ip_idx]
def get_port_by_idx(self, port_idx: int) -> str:
@@ -916,6 +949,13 @@ class ActionManager:
:return: The port.
:rtype: str
"""
if port_idx >= len(self.ports):
msg = (
f"Error: agent attempted to perform an action on port {port_idx} but this"
f" is out of range for its action space. Port list: {self.ip_address_list}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.ports[port_idx]
def get_nic_num_by_idx(self, node_idx: int, nic_idx: int) -> int:
@@ -958,6 +998,12 @@ class ActionManager:
:return: The constructed ActionManager.
:rtype: ActionManager
"""
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
ip_address_order = cfg["options"].pop("ip_address_order", {})
ip_address_list = []
for entry in ip_address_order:
@@ -967,13 +1013,22 @@ class ActionManager:
ip_address = node_obj.network_interface[nic_num].ip_address
ip_address_list.append(ip_address)
if not ip_address_list:
node_names = [n["node_name"] for n in cfg.get("nodes", {})]
for node_name in node_names:
node_obj = game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
ip_address_list.append(nic_obj.ip_address)
obj = cls(
game=game,
actions=cfg["action_list"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,
ip_address_list=ip_address_list or None,
ip_address_list=ip_address_list,
act_map=cfg.get("action_map"),
)

52
src/primaite/game/agent/data_manipulation_bot.py
View File

@@ -1,52 +0,0 @@
import random
from typing import Dict, List, Tuple
from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
class DataManipulationAgent(AbstractScriptedAgent):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
data_manipulation_bots: List["DataManipulationBot"] = []
next_execution_timestep: int = 0
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
def _set_next_execution_timestep(self, timestep: int) -> None:
"""Set the next execution timestep with a configured random variance.
:param timestep: The timestep to add variance to.
"""
random_timestep_increment = random.randint(
-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance
)
self.next_execution_timestep = timestep + random_timestep_increment
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
"""Randomly sample an action from the action space.
:param obs: _description_
:type obs: ObsType
:param reward: _description_, defaults to None
:type reward: float, optional
:return: _description_
:rtype: Tuple[str, Dict]
"""
current_timestep = self.action_manager.game.step_counter
if current_timestep < self.next_execution_timestep:
return "DONOTHING", {"dummy": 0}
self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency)
return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
def reset_agent_for_episode(self) -> None:
"""Set the next execution timestep when the episode resets."""
super().reset_agent_for_episode()
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)

72
src/primaite/game/agent/interface.py
View File

@@ -1,18 +1,38 @@
"""Interface for agents."""
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium.core import ActType, ObsType
from pydantic import BaseModel, model_validator
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.observations import ObservationManager
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.interface.request import RequestFormat, RequestResponse
if TYPE_CHECKING:
pass
class AgentActionHistoryItem(BaseModel):
"""One entry of an agent's action log - what the agent did and how the simulator responded in 1 step."""
timestep: int
"""Timestep of this action."""
action: str
"""CAOS Action name."""
parameters: Dict[str, Any]
"""CAOS parameters for the given action."""
request: RequestFormat
"""The request that was sent to the simulation based on the CAOS action chosen."""
response: RequestResponse
"""The response sent back by the simulator for this action."""
class AgentStartSettings(BaseModel):
"""Configuration values for when an agent starts performing actions."""
@@ -90,6 +110,7 @@ class AbstractAgent(ABC):
self.observation_manager: Optional[ObservationManager] = observation_space
self.reward_function: Optional[RewardFunction] = reward_function
self.agent_settings = agent_settings or AgentSettings()
self.action_history: List[AgentActionHistoryItem] = []
def update_observation(self, state: Dict) -> ObsType:
"""
@@ -109,10 +130,10 @@ class AbstractAgent(ABC):
:return: Reward from the state.
:rtype: float
"""
return self.reward_function.update(state)
return self.reward_function.update(state=state, last_action_response=self.action_history[-1])
@abstractmethod
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""
Return an action to be taken in the environment.
@@ -120,8 +141,8 @@ class AbstractAgent(ABC):
:param obs: Observation of the environment.
:type obs: ObsType
:param reward: Reward from the previous action, defaults to None TODO: should this parameter even be accepted?
:type reward: float, optional
:param timestep: The current timestep in the simulation, used for non-RL agents. Optional
:type timestep: int
:return: Action to be taken in the environment.
:rtype: Tuple[str, Dict]
"""
@@ -136,31 +157,24 @@ class AbstractAgent(ABC):
request = self.action_manager.form_request(action_identifier=action, action_options=options)
return request
def reset_agent_for_episode(self) -> None:
"""Agent reset logic should go here."""
pass
def process_action_response(
self, timestep: int, action: str, parameters: Dict[str, Any], request: RequestFormat, response: RequestResponse
) -> None:
"""Process the response from the most recent action."""
self.action_history.append(
AgentActionHistoryItem(
timestep=timestep, action=action, parameters=parameters, request=request, response=response
)
)
class AbstractScriptedAgent(AbstractAgent):
"""Base class for actors which generate their own behaviour."""
...
class RandomAgent(AbstractScriptedAgent):
"""Agent that ignores its observation and acts completely at random."""
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
"""Randomly sample an action from the action space.
:param obs: _description_
:type obs: ObsType
:param reward: _description_, defaults to None
:type reward: float, optional
:return: _description_
:rtype: Tuple[str, Dict]
"""
return self.action_manager.get_action(self.action_manager.space.sample())
@abstractmethod
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""Return an action to be taken in the environment."""
return super().get_action(obs=obs, timestep=timestep)
class ProxyAgent(AbstractAgent):
@@ -183,14 +197,14 @@ class ProxyAgent(AbstractAgent):
self.most_recent_action: ActType
self.flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""
Return the agent's most recent action, formatted in CAOS format.
:param obs: Observation for the agent. Not used by ProxyAgents, but required by the interface.
:type obs: ObsType
:param reward: Reward value for the agent. Not used by ProxyAgents, defaults to None.
:type reward: float, optional
:param timestep: Current simulation timestep. Not used by ProxyAgents, bur required for the interface.
:type timestep: int
:return: Action to be taken in CAOS format.
:rtype: Tuple[str, Dict]
"""

1010
src/primaite/game/agent/observations.py
View File

File diff suppressed because it is too large Load Diff

0
src/primaite/game/agent/observations/__init__.py Normal file
View File

188
src/primaite/game/agent/observations/agent_observations.py Normal file
View File

@@ -0,0 +1,188 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite.game.agent.observations.node_observations import NodeObservation
from primaite.game.agent.observations.observations import (
AbstractObservation,
AclObservation,
ICSObservation,
LinkObservation,
NullObservation,
)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class UC2BlueObservation(AbstractObservation):
"""Container for all observations used by the blue agent in UC2.
TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler
for the purpose of compiling several observation components.
"""
def __init__(
self,
nodes: List[NodeObservation],
links: List[LinkObservation],
acl: AclObservation,
ics: ICSObservation,
where: Optional[List[str]] = None,
) -> None:
"""Initialise UC2 blue observation.
:param nodes: List of node observations
:type nodes: List[NodeObservation]
:param links: List of link observations
:type links: List[LinkObservation]
:param acl: The Access Control List observation
:type acl: AclObservation
:param ics: The ICS observation
:type ics: ICSObservation
:param where: Where in the simulation state dict to find information. Not used in this particular observation
because it only compiles other observations and doesn't contribute any new information, defaults to None
:type where: Optional[List[str]], optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.nodes: List[NodeObservation] = nodes
self.links: List[LinkObservation] = links
self.acl: AclObservation = acl
self.ics: ICSObservation = ics
self.default_observation: Dict = {
"NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
"LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)},
"ACL": self.acl.default_observation,
"ICS": self.ics.default_observation,
}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
obs = {}
obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)}
obs["ACL"] = self.acl.observe(state)
obs["ICS"] = self.ics.observe(state)
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Space
:rtype: spaces.Space
"""
return spaces.Dict(
{
"NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
"LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}),
"ACL": self.acl.space,
"ICS": self.ics.space,
}
)
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation":
"""Create UC2 blue observation from a config.
:param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes,
links, ACL and ICS observations.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:return: Constructed UC2 blue observation
:rtype: UC2BlueObservation
"""
node_configs = config["nodes"]
num_services_per_node = config["num_services_per_node"]
num_folders_per_node = config["num_folders_per_node"]
num_files_per_folder = config["num_files_per_folder"]
num_nics_per_node = config["num_nics_per_node"]
nodes = [
NodeObservation.from_config(
config=n,
game=game,
num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder,
num_nics_per_node=num_nics_per_node,
)
for n in node_configs
]
link_configs = config["links"]
links = [LinkObservation.from_config(config=link, game=game) for link in link_configs]
acl_config = config["acl"]
acl = AclObservation.from_config(config=acl_config, game=game)
ics_config = config["ics"]
ics = ICSObservation.from_config(config=ics_config, game=game)
new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"])
return new
class UC2RedObservation(AbstractObservation):
"""Container for all observations used by the red agent in UC2."""
def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None:
super().__init__()
self.where: Optional[List[str]] = where
self.nodes: List[NodeObservation] = nodes
self.default_observation: Dict = {
"NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation."""
if self.where is None:
return self.default_observation
obs = {}
obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
return obs
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
return spaces.Dict(
{
"NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
}
)
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2RedObservation":
"""
Create UC2 red observation from a config.
:param config: Dictionary containing the configuration for this UC2 red observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
"""
node_configs = config["nodes"]
nodes = [NodeObservation.from_config(config=cfg, game=game) for cfg in node_configs]
return cls(nodes=nodes, where=["network"])
class UC2GreenObservation(NullObservation):
"""Green agent observation. As the green agent's actions don't depend on the observation, this is empty."""
pass

177
src/primaite/game/agent/observations/file_system_observations.py Normal file
View File

@@ -0,0 +1,177 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class FileObservation(AbstractObservation):
"""Observation of a file on a node in the network."""
def __init__(self, where: Optional[Tuple[str]] = None) -> None:
"""
Initialise file observation.
:param where: Store information about where in the simulation state dictionary to find the relevant information.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a file looks like this:
['network','nodes',<node_hostname>,'file_system', 'folders',<folder_name>,'files',<file_name>]
:type where: Optional[List[str]]
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.default_observation: spaces.Space = {"health_status": 0}
"Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted."
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
file_state = access_from_nested_dict(state, self.where)
if file_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {"health_status": file_state["visible_status"]}
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.
:return: Gymnasium space
:rtype: spaces.Space
"""
return spaces.Dict({"health_status": spaces.Discrete(6)})
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame", parent_where: List[str] = None) -> "FileObservation":
"""Create file observation from a config.
:param config: Dictionary containing the configuration for this file observation.
:type config: Dict
:param game: _description_
:type game: PrimaiteGame
:param parent_where: _description_, defaults to None
:type parent_where: _type_, optional
:return: _description_
:rtype: _type_
"""
return cls(where=parent_where + ["files", config["file_name"]])
class FolderObservation(AbstractObservation):
"""Folder observation, including files inside of the folder."""
def __init__(
self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder: int = 2
) -> None:
"""Initialise folder Observation, including files inside the folder.
:param where: Where in the simulation state dictionary to find the relevant information for this folder.
A typical location for a file looks like this:
['network','nodes',<node_hostname>,'file_system', 'folders',<folder_name>]
:type where: Optional[List[str]]
:param max_files: As size of the space must remain static, define max files that can be in this folder
, defaults to 5
:type max_files: int, optional
:param file_positions: Defines the positioning within the observation space of particular files. This ensures
that even if new files are created, the existing files will always occupy the same space in the observation
space. The keys must be between 1 and max_files. Providing file_positions will reserve a spot in the
observation space for a file with that name, even if it's temporarily deleted, if it reappears with the same
name, it will take the position defined in this dict. Defaults to {}
:type file_positions: Dict[int, str], optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.files: List[FileObservation] = files
while len(self.files) < num_files_per_folder:
self.files.append(FileObservation())
while len(self.files) > num_files_per_folder:
truncated_file = self.files.pop()
msg = f"Too many files in folder observation. Truncating file {truncated_file}"
_LOGGER.warning(msg)
self.default_observation = {
"health_status": 0,
"FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)},
}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
folder_state = access_from_nested_dict(state, self.where)
if folder_state is NOT_PRESENT_IN_STATE:
return self.default_observation
health_status = folder_state["health_status"]
obs = {}
obs["health_status"] = health_status
obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)}
return obs
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.
:return: Gymnasium space
:rtype: spaces.Space
"""
return spaces.Dict(
{
"health_status": spaces.Discrete(6),
"FILES": spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)}),
}
)
@classmethod
def from_config(
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]], num_files_per_folder: int = 2
) -> "FolderObservation":
"""Create folder observation from a config. Also creates child file observations.
:param config: Dictionary containing the configuration for this folder observation. Includes the name of the
folder and the files inside of it.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary to find the information about this folder's
parent node. A typical location for a node ``where`` can be:
['network','nodes',<node_hostname>,'file_system']
:type parent_where: Optional[List[str]]
:param num_files_per_folder: How many spaces for files are in this folder observation (to preserve static
observation size) , defaults to 2
:type num_files_per_folder: int, optional
:return: Constructed folder observation
:rtype: FolderObservation
"""
where = parent_where + ["folders", config["folder_name"]]
file_configs = config["files"]
files = [FileObservation.from_config(config=f, game=game, parent_where=where) for f in file_configs]
return cls(where=where, files=files, num_files_per_folder=num_files_per_folder)

188
src/primaite/game/agent/observations/nic_observations.py Normal file
View File

@@ -0,0 +1,188 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite.game.agent.observations.observations import AbstractObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.nmne import CAPTURE_NMNE
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class NicObservation(AbstractObservation):
"""Observation of a Network Interface Card (NIC) in the network."""
low_nmne_threshold: int = 0
"""The minimum number of malicious network events to be considered low."""
med_nmne_threshold: int = 5
"""The minimum number of malicious network events to be considered medium."""
high_nmne_threshold: int = 10
"""The minimum number of malicious network events to be considered high."""
global CAPTURE_NMNE
@property
def default_observation(self) -> Dict:
"""The default NIC observation dict."""
data = {"nic_status": 0}
if CAPTURE_NMNE:
data.update({"NMNE": {"inbound": 0, "outbound": 0}})
return data
def __init__(
self,
where: Optional[Tuple[str]] = None,
low_nmne_threshold: Optional[int] = 0,
med_nmne_threshold: Optional[int] = 5,
high_nmne_threshold: Optional[int] = 10,
) -> None:
"""Initialise NIC observation.
:param where: Where in the simulation state dictionary to find the relevant information for this NIC. A typical
example may look like this:
['network','nodes',<node_hostname>,'NICs',<nic_number>]
If None, this denotes that the NIC does not exist and the observation will be populated with zeroes.
:type where: Optional[Tuple[str]], optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
global CAPTURE_NMNE
if CAPTURE_NMNE:
self.nmne_inbound_last_step: int = 0
"""NMNEs persist for the whole episode, but we want to count per step. Keeping track of last step count lets
us find the difference."""
self.nmne_outbound_last_step: int = 0
"""NMNEs persist for the whole episode, but we want to count per step. Keeping track of last step count lets
us find the difference."""
if low_nmne_threshold or med_nmne_threshold or high_nmne_threshold:
self._validate_nmne_categories(
low_nmne_threshold=low_nmne_threshold,
med_nmne_threshold=med_nmne_threshold,
high_nmne_threshold=high_nmne_threshold,
)
def _validate_nmne_categories(
self, low_nmne_threshold: int = 0, med_nmne_threshold: int = 5, high_nmne_threshold: int = 10
):
"""
Validates the nmne threshold config.
If the configuration is valid, the thresholds will be set, otherwise, an exception is raised.
:param: low_nmne_threshold: The minimum number of malicious network events to be considered low
:param: med_nmne_threshold: The minimum number of malicious network events to be considered medium
:param: high_nmne_threshold: The minimum number of malicious network events to be considered high
"""
if high_nmne_threshold <= med_nmne_threshold:
raise Exception(
f"nmne_categories: high nmne count ({high_nmne_threshold}) must be greater "
f"than medium nmne count ({med_nmne_threshold})"
)
if med_nmne_threshold <= low_nmne_threshold:
raise Exception(
f"nmne_categories: medium nmne count ({med_nmne_threshold}) must be greater "
f"than low nmne count ({low_nmne_threshold})"
)
self.high_nmne_threshold = high_nmne_threshold
self.med_nmne_threshold = med_nmne_threshold
self.low_nmne_threshold = low_nmne_threshold
def _categorise_mne_count(self, nmne_count: int) -> int:
"""
Categorise the number of Malicious Network Events (NMNEs) into discrete bins.
This helps in classifying the severity or volume of MNEs into manageable levels for the agent.
Bins are defined as follows:
- 0: No MNEs detected (0 events).
- 1: Low number of MNEs (default 1-5 events).
- 2: Moderate number of MNEs (default 6-10 events).
- 3: High number of MNEs (default more than 10 events).
:param nmne_count: Number of MNEs detected.
:return: Bin number corresponding to the number of MNEs. Returns 0, 1, 2, or 3 based on the detected MNE count.
"""
if nmne_count > self.high_nmne_threshold:
return 3
elif nmne_count > self.med_nmne_threshold:
return 2
elif nmne_count > self.low_nmne_threshold:
return 1
return 0
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
nic_state = access_from_nested_dict(state, self.where)
if nic_state is NOT_PRESENT_IN_STATE:
return self.default_observation
else:
obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2}
if CAPTURE_NMNE:
obs_dict.update({"NMNE": {}})
direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
inbound_count = inbound_keywords.get("*", 0)
outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {})
outbound_count = outbound_keywords.get("*", 0)
obs_dict["NMNE"]["inbound"] = self._categorise_mne_count(inbound_count - self.nmne_inbound_last_step)
obs_dict["NMNE"]["outbound"] = self._categorise_mne_count(outbound_count - self.nmne_outbound_last_step)
self.nmne_inbound_last_step = inbound_count
self.nmne_outbound_last_step = outbound_count
return obs_dict
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
space = spaces.Dict({"nic_status": spaces.Discrete(3)})
if CAPTURE_NMNE:
space["NMNE"] = spaces.Dict({"inbound": spaces.Discrete(4), "outbound": spaces.Discrete(4)})
return space
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]]) -> "NicObservation":
"""Create NIC observation from a config.
:param config: Dictionary containing the configuration for this NIC observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary to find the information about this NIC's parent
node. A typical location for a node ``where`` can be: ['network','nodes',<node_hostname>]
:type parent_where: Optional[List[str]]
:return: Constructed NIC observation
:rtype: NicObservation
"""
low_nmne_threshold = None
med_nmne_threshold = None
high_nmne_threshold = None
if game and game.options and game.options.thresholds and game.options.thresholds.get("nmne"):
threshold = game.options.thresholds["nmne"]
low_nmne_threshold = int(threshold.get("low")) if threshold.get("low") is not None else None
med_nmne_threshold = int(threshold.get("medium")) if threshold.get("medium") is not None else None
high_nmne_threshold = int(threshold.get("high")) if threshold.get("high") is not None else None
return cls(
where=parent_where + ["NICs", config["nic_num"]],
low_nmne_threshold=low_nmne_threshold,
med_nmne_threshold=med_nmne_threshold,
high_nmne_threshold=high_nmne_threshold,
)

200
src/primaite/game/agent/observations/node_observations.py Normal file
View File

@@ -0,0 +1,200 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite import getLogger
from primaite.game.agent.observations.file_system_observations import FolderObservation
from primaite.game.agent.observations.nic_observations import NicObservation
from primaite.game.agent.observations.observations import AbstractObservation
from primaite.game.agent.observations.software_observation import ServiceObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class NodeObservation(AbstractObservation):
"""Observation of a node in the network. Includes services, folders and NICs."""
def __init__(
self,
where: Optional[Tuple[str]] = None,
services: List[ServiceObservation] = [],
folders: List[FolderObservation] = [],
network_interfaces: List[NicObservation] = [],
logon_status: bool = False,
num_services_per_node: int = 2,
num_folders_per_node: int = 2,
num_files_per_folder: int = 2,
num_nics_per_node: int = 2,
) -> None:
"""
Configurable observation for a node in the simulation.
:param where: Where in the simulation state dictionary for find relevant information for this observation.
A typical location for a node looks like this:
['network','nodes',<hostname>]. If empty list, a default null observation will be output, defaults to []
:type where: List[str], optional
:param services: Mapping between position in observation space and service name, defaults to {}
:type services: Dict[int,str], optional
:param max_services: Max number of services that can be presented in observation space for this node
, defaults to 2
:type max_services: int, optional
:param folders: Mapping between position in observation space and folder name, defaults to {}
:type folders: Dict[int,str], optional
:param max_folders: Max number of folders in this node's obs space, defaults to 2
:type max_folders: int, optional
:param network_interfaces: Mapping between position in observation space and NIC idx, defaults to {}
:type network_interfaces: Dict[int,str], optional
:param max_nics: Max number of network interfaces in this node's obs space, defaults to 5
:type max_nics: int, optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.services: List[ServiceObservation] = services
while len(self.services) < num_services_per_node:
# add empty service observation without `where` parameter so it always returns default (blank) observation
self.services.append(ServiceObservation())
while len(self.services) > num_services_per_node:
truncated_service = self.services.pop()
msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
_LOGGER.warning(msg)
# truncate service list
self.folders: List[FolderObservation] = folders
# add empty folder observation without `where` parameter that will always return default (blank) observations
while len(self.folders) < num_folders_per_node:
self.folders.append(FolderObservation(num_files_per_folder=num_files_per_folder))
while len(self.folders) > num_folders_per_node:
truncated_folder = self.folders.pop()
msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}"
_LOGGER.warning(msg)
self.network_interfaces: List[NicObservation] = network_interfaces
while len(self.network_interfaces) < num_nics_per_node:
self.network_interfaces.append(NicObservation())
while len(self.network_interfaces) > num_nics_per_node:
truncated_nic = self.network_interfaces.pop()
msg = f"Too many NICs in Node observation for node. Truncating service {truncated_nic.where[-1]}"
_LOGGER.warning(msg)
self.logon_status: bool = logon_status
self.default_observation: Dict = {
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
"FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)},
"NICS": {i + 1: n.default_observation for i, n in enumerate(self.network_interfaces)},
"operating_status": 0,
}
if self.logon_status:
self.default_observation["logon_status"] = 0
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
node_state = access_from_nested_dict(state, self.where)
if node_state is NOT_PRESENT_IN_STATE:
return self.default_observation
obs = {}
obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
obs["operating_status"] = node_state["operating_state"]
obs["NICS"] = {
i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces)
}
if self.logon_status:
obs["logon_status"] = 0
return obs
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
space_shape = {
"SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
"FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
"operating_status": spaces.Discrete(5),
"NICS": spaces.Dict(
{i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}
),
}
if self.logon_status:
space_shape["logon_status"] = spaces.Discrete(3)
return spaces.Dict(space_shape)
@classmethod
def from_config(
cls,
config: Dict,
game: "PrimaiteGame",
parent_where: Optional[List[str]] = None,
num_services_per_node: int = 2,
num_folders_per_node: int = 2,
num_files_per_folder: int = 2,
num_nics_per_node: int = 2,
) -> "NodeObservation":
"""Create node observation from a config. Also creates child service, folder and NIC observations.
:param config: Dictionary containing the configuration for this node observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary to find the information about this node's parent
network. A typical location for it would be: ['network',]
:type parent_where: Optional[List[str]]
:param num_services_per_node: How many spaces for services are in this node observation (to preserve static
observation size) , defaults to 2
:type num_services_per_node: int, optional
:param num_folders_per_node: How many spaces for folders are in this node observation (to preserve static
observation size) , defaults to 2
:type num_folders_per_node: int, optional
:param num_files_per_folder: How many spaces for files are in the folder observations (to preserve static
observation size) , defaults to 2
:type num_files_per_folder: int, optional
:return: Constructed node observation
:rtype: NodeObservation
"""
node_hostname = config["node_hostname"]
if parent_where is None:
where = ["network", "nodes", node_hostname]
else:
where = parent_where + ["nodes", node_hostname]
svc_configs = config.get("services", {})
services = [ServiceObservation.from_config(config=c, game=game, parent_where=where) for c in svc_configs]
folder_configs = config.get("folders", {})
folders = [
FolderObservation.from_config(
config=c, game=game, parent_where=where + ["file_system"], num_files_per_folder=num_files_per_folder
)
for c in folder_configs
]
# create some configs for the NIC observation in the format {"nic_num":1}, {"nic_num":2}, {"nic_num":3}, etc.
nic_configs = [{"nic_num": i for i in range(num_nics_per_node)}]
network_interfaces = [NicObservation.from_config(config=c, game=game, parent_where=where) for c in nic_configs]
logon_status = config.get("logon_status", False)
return cls(
where=where,
services=services,
folders=folders,
network_interfaces=network_interfaces,
logon_status=logon_status,
num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder,
num_nics_per_node=num_nics_per_node,
)

73
src/primaite/game/agent/observations/observation_manager.py Normal file
View File

@@ -0,0 +1,73 @@
from typing import Dict, TYPE_CHECKING
from gymnasium.core import ObsType
from primaite.game.agent.observations.agent_observations import (
UC2BlueObservation,
UC2GreenObservation,
UC2RedObservation,
)
from primaite.game.agent.observations.observations import AbstractObservation
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class ObservationManager:
"""
Manage the observations of an Agent.
The observation space has the purpose of:
1. Reading the outputted state from the PrimAITE Simulation.
2. Selecting parts of the simulation state that are requested by the simulation config
3. Formatting this information so an agent can use it to make decisions.
"""
# TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed
# to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next
# refactor.
def __init__(self, observation: AbstractObservation) -> None:
"""Initialise observation space.
:param observation: Observation object
:type observation: AbstractObservation
"""
self.obs: AbstractObservation = observation
self.current_observation: ObsType
def update(self, state: Dict) -> Dict:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
"""
self.current_observation = self.obs.observe(state)
return self.current_observation
@property
def space(self) -> None:
"""Gymnasium space object describing the observation space shape."""
return self.obs.space
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "ObservationManager":
"""Create observation space from a config.
:param config: Dictionary containing the configuration for this observation space.
It should contain the key 'type' which selects which observation class to use (from a choice of:
UC2BlueObservation, UC2RedObservation, UC2GreenObservation)
The other key is 'options' which are passed to the constructor of the selected observation class.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
"""
if config["type"] == "UC2BlueObservation":
return cls(UC2BlueObservation.from_config(config.get("options", {}), game=game))
elif config["type"] == "UC2RedObservation":
return cls(UC2RedObservation.from_config(config.get("options", {}), game=game))
elif config["type"] == "UC2GreenObservation":
return cls(UC2GreenObservation.from_config(config.get("options", {}), game=game))
else:
raise ValueError("Observation space type invalid")

309
src/primaite/game/agent/observations/observations.py Normal file
View File

@@ -0,0 +1,309 @@
"""Manages the observation space for the agent."""
from abc import ABC, abstractmethod
from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class AbstractObservation(ABC):
"""Abstract class for an observation space component."""
@abstractmethod
def observe(self, state: Dict) -> Any:
"""
Return an observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Any
"""
pass
@property
@abstractmethod
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space."""
pass
@classmethod
@abstractmethod
def from_config(cls, config: Dict, game: "PrimaiteGame"):
"""Create this observation space component form a serialised format.
The `game` parameter is for a the PrimaiteGame object that spawns this component.
"""
pass
class LinkObservation(AbstractObservation):
"""Observation of a link in the network."""
default_observation: spaces.Space = {"PROTOCOLS": {"ALL": 0}}
"Default observation is what should be returned when the link doesn't exist."
def __init__(self, where: Optional[Tuple[str]] = None) -> None:
"""Initialise link observation.
:param where: Store information about where in the simulation state dictionary to find the relevant information.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a service looks like this:
`['network','nodes',<node_hostname>,'servics', <service_name>]`
:type where: Optional[List[str]]
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
link_state = access_from_nested_dict(state, self.where)
if link_state is NOT_PRESENT_IN_STATE:
return self.default_observation
bandwidth = link_state["bandwidth"]
load = link_state["current_load"]
if load == 0:
utilisation_category = 0
else:
utilisation_fraction = load / bandwidth
# 0 is UNUSED, 1 is 0%-10%. 2 is 10%-20%. 3 is 20%-30%. And so on... 10 is exactly 100%
utilisation_category = int(utilisation_fraction * 9) + 1
# TODO: once the links support separte load per protocol, this needs amendment to reflect that.
return {"PROTOCOLS": {"ALL": min(utilisation_category, 10)}}
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.
:return: Gymnasium space
:rtype: spaces.Space
"""
return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "LinkObservation":
"""Create link observation from a config.
:param config: Dictionary containing the configuration for this link observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:return: Constructed link observation
:rtype: LinkObservation
"""
return cls(where=["network", "links", game.ref_map_links[config["link_ref"]]])
class AclObservation(AbstractObservation):
"""Observation of an Access Control List (ACL) in the network."""
# TODO: should where be optional, and we can use where=None to pad the observation space?
# definitely the current approach does not support tracking files that aren't specified by name, for example
# if a file is created at runtime, we have currently got no way of telling the observation space to track it.
# this needs adding, but not for the MVP.
def __init__(
self,
node_ip_to_id: Dict[str, int],
ports: List[int],
protocols: List[str],
where: Optional[Tuple[str]] = None,
num_rules: int = 10,
) -> None:
"""Initialise ACL observation.
:param node_ip_to_id: Mapping between IP address and ID.
:type node_ip_to_id: Dict[str, int]
:param ports: List of ports which are part of the game that define the ordering when converting to an ID
:type ports: List[int]
:param protocols: List of protocols which are part of the game, defines ordering when converting to an ID
:type protocols: list[str]
:param where: Where in the simulation state dictionary to find the relevant information for this ACL. A typical
example may look like this:
['network','nodes',<router_hostname>,'acl','acl']
:type where: Optional[Tuple[str]], optional
:param num_rules: , defaults to 10
:type num_rules: int, optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.num_rules: int = num_rules
self.node_to_id: Dict[str, int] = node_ip_to_id
"List of node IP addresses, order in this list determines how they are converted to an ID"
self.port_to_id: Dict[int, int] = {port: i + 2 for i, port in enumerate(ports)}
"List of ports which are part of the game that define the ordering when converting to an ID"
self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)}
"List of protocols which are part of the game, defines ordering when converting to an ID"
self.default_observation: Dict = {
i
+ 1: {
"position": i,
"permission": 0,
"source_node_id": 0,
"source_port": 0,
"dest_node_id": 0,
"dest_port": 0,
"protocol": 0,
}
for i in range(self.num_rules)
}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
acl_state: Dict = access_from_nested_dict(state, self.where)
if acl_state is NOT_PRESENT_IN_STATE:
return self.default_observation
# TODO: what if the ACL has more rules than num of max rules for obs space
obs = {}
acl_items = dict(acl_state.items())
i = 1 # don't show rule 0 for compatibility reasons.
while i < self.num_rules + 1:
rule_state = acl_items[i]
if rule_state is None:
obs[i] = {
"position": i - 1,
"permission": 0,
"source_node_id": 0,
"source_port": 0,
"dest_node_id": 0,
"dest_port": 0,
"protocol": 0,
}
else:
src_ip = rule_state["src_ip_address"]
src_node_id = 1 if src_ip is None else self.node_to_id[IPv4Address(src_ip)]
dst_ip = rule_state["dst_ip_address"]
dst_node_ip = 1 if dst_ip is None else self.node_to_id[IPv4Address(dst_ip)]
src_port = rule_state["src_port"]
src_port_id = 1 if src_port is None else self.port_to_id[src_port]
dst_port = rule_state["dst_port"]
dst_port_id = 1 if dst_port is None else self.port_to_id[dst_port]
protocol = rule_state["protocol"]
protocol_id = 1 if protocol is None else self.protocol_to_id[protocol]
obs[i] = {
"position": i - 1,
"permission": rule_state["action"],
"source_node_id": src_node_id,
"source_port": src_port_id,
"dest_node_id": dst_node_ip,
"dest_port": dst_port_id,
"protocol": protocol_id,
}
i += 1
return obs
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.
:return: Gymnasium space
:rtype: spaces.Space
"""
return spaces.Dict(
{
i
+ 1: spaces.Dict(
{
"position": spaces.Discrete(self.num_rules),
"permission": spaces.Discrete(3),
# adding two to lengths is to account for reserved values 0 (unused) and 1 (any)
"source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
"source_port": spaces.Discrete(len(self.port_to_id) + 2),
"dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
"dest_port": spaces.Discrete(len(self.port_to_id) + 2),
"protocol": spaces.Discrete(len(self.protocol_to_id) + 2),
}
)
for i in range(self.num_rules)
}
)
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "AclObservation":
"""Generate ACL observation from a config.
:param config: Dictionary containing the configuration for this ACL observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:return: Observation object
:rtype: AclObservation
"""
max_acl_rules = config["options"]["max_acl_rules"]
node_ip_to_idx = {}
for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
node_ref = ip_map_config["node_hostname"]
nic_num = ip_map_config["nic_num"]
node_obj = game.simulation.network.nodes[game.ref_map_nodes[node_ref]]
nic_obj = node_obj.network_interface[nic_num]
node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2
router_hostname = config["router_hostname"]
return cls(
node_ip_to_id=node_ip_to_idx,
ports=game.options.ports,
protocols=game.options.protocols,
where=["network", "nodes", router_hostname, "acl", "acl"],
num_rules=max_acl_rules,
)
class NullObservation(AbstractObservation):
"""Null observation, returns a single 0 value for the observation space."""
def __init__(self, where: Optional[List[str]] = None):
"""Initialise null observation."""
self.default_observation: Dict = {}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation."""
return 0
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
return spaces.Discrete(1)
@classmethod
def from_config(cls, config: Dict, game: Optional["PrimaiteGame"] = None) -> "NullObservation":
"""
Create null observation from a config.
The parameters are ignored, they are here to match the signature of the other observation classes.
"""
return cls()
class ICSObservation(NullObservation):
"""ICS observation placeholder, currently not implemented so always returns a single 0."""
pass

163
src/primaite/game/agent/observations/software_observation.py Normal file
View File

@@ -0,0 +1,163 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite.game.agent.observations.observations import AbstractObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class ServiceObservation(AbstractObservation):
"""Observation of a service in the network."""
default_observation: spaces.Space = {"operating_status": 0, "health_status": 0}
"Default observation is what should be returned when the service doesn't exist."
def __init__(self, where: Optional[Tuple[str]] = None) -> None:
"""Initialise service observation.
:param where: Store information about where in the simulation state dictionary to find the relevant information.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a service looks like this:
`['network','nodes',<node_hostname>,'services', <service_name>]`
:type where: Optional[List[str]]
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
service_state = access_from_nested_dict(state, self.where)
if service_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {
"operating_status": service_state["operating_state"],
"health_status": service_state["health_state_visible"],
}
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
@classmethod
def from_config(
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]] = None
) -> "ServiceObservation":
"""Create service observation from a config.
:param config: Dictionary containing the configuration for this service observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional.
:type parent_where: Optional[List[str]], optional
:return: Constructed service observation
:rtype: ServiceObservation
"""
return cls(where=parent_where + ["services", config["service_name"]])
class ApplicationObservation(AbstractObservation):
"""Observation of an application in the network."""
default_observation: spaces.Space = {"operating_status": 0, "health_status": 0, "num_executions": 0}
"Default observation is what should be returned when the application doesn't exist."
def __init__(self, where: Optional[Tuple[str]] = None) -> None:
"""Initialise application observation.
:param where: Store information about where in the simulation state dictionary to find the relevant information.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a service looks like this:
`['network','nodes',<node_hostname>,'applications', <application_name>]`
:type where: Optional[List[str]]
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
app_state = access_from_nested_dict(state, self.where)
if app_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {
"operating_status": app_state["operating_state"],
"health_status": app_state["health_state_visible"],
"num_executions": self._categorise_num_executions(app_state["num_executions"]),
}
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
return spaces.Dict(
{
"operating_status": spaces.Discrete(7),
"health_status": spaces.Discrete(6),
"num_executions": spaces.Discrete(4),
}
)
@classmethod
def from_config(
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]] = None
) -> "ApplicationObservation":
"""Create application observation from a config.
:param config: Dictionary containing the configuration for this service observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional.
:type parent_where: Optional[List[str]], optional
:return: Constructed service observation
:rtype: ApplicationObservation
"""
return cls(where=parent_where + ["services", config["application_name"]])
@classmethod
def _categorise_num_executions(cls, num_executions: int) -> int:
"""
Categorise the number of executions of an application.
Helps classify the number of application executions into different categories.
Current categories:
- 0: Application is never executed
- 1: Application is executed a low number of times (1-5)
- 2: Application is executed often (6-10)
- 3: Application is executed a high number of times (more than 10)
:param: num_executions: Number of times the application is executed
"""
if num_executions > 10:
return 3
elif num_executions > 5:
return 2
elif num_executions > 0:
return 1
return 0

158
src/primaite/game/agent/rewards.py
View File

@@ -13,7 +13,7 @@ the structure:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_name: database_server
folder_name: database
file_name: database.db
@@ -21,16 +21,21 @@ the structure:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
node_name: web_server
service_ref: web_server_database_client
```
"""
from abc import abstractmethod
from typing import Dict, List, Tuple, Type
from typing import Callable, Dict, List, Optional, Tuple, Type, TYPE_CHECKING
from typing_extensions import Never
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
if TYPE_CHECKING:
from primaite.game.agent.interface import AgentActionHistoryItem
_LOGGER = getLogger(__name__)
@@ -38,7 +43,7 @@ class AbstractReward:
"""Base class for reward function components."""
@abstractmethod
def calculate(self, state: Dict) -> float:
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Calculate the reward for the current state."""
return 0.0
@@ -58,7 +63,7 @@ class AbstractReward:
class DummyReward(AbstractReward):
"""Dummy reward function component which always returns 0."""
def calculate(self, state: Dict) -> float:
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Calculate the reward for the current state."""
return 0.0
@@ -98,7 +103,7 @@ class DatabaseFileIntegrity(AbstractReward):
file_name,
]
def calculate(self, state: Dict) -> float:
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Calculate the reward for the current state.
:param state: The current state of the simulation.
@@ -106,7 +111,7 @@ class DatabaseFileIntegrity(AbstractReward):
"""
database_file_state = access_from_nested_dict(state, self.location_in_state)
if database_file_state is NOT_PRESENT_IN_STATE:
_LOGGER.info(
_LOGGER.debug(
f"Could not calculate {self.__class__} reward because "
"simulation state did not contain enough information."
)
@@ -153,7 +158,7 @@ class WebServer404Penalty(AbstractReward):
"""
self.location_in_state = ["network", "nodes", node_hostname, "services", service_name]
def calculate(self, state: Dict) -> float:
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Calculate the reward for the current state.
:param state: The current state of the simulation.
@@ -184,7 +189,7 @@ class WebServer404Penalty(AbstractReward):
service_name = config.get("service_name")
if not (node_hostname and service_name):
msg = (
f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not "
f"{cls.__name__} could not be initialised from config because node_name and service_ref were not "
"found in reward config."
)
_LOGGER.warning(msg)
@@ -203,19 +208,30 @@ class WebpageUnavailablePenalty(AbstractReward):
:param node_hostname: Hostname of the node which has the web browser.
:type node_hostname: str
"""
self._node = node_hostname
self.location_in_state = ["network", "nodes", node_hostname, "applications", "WebBrowser"]
self._node: str = node_hostname
self.location_in_state: List[str] = ["network", "nodes", node_hostname, "applications", "WebBrowser"]
self._last_request_failed: bool = False
def calculate(self, state: Dict) -> float:
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""
Calculate the reward based on current simulation state.
Calculate the reward based on current simulation state, and the recent agent action.
:param state: The current state of the simulation.
:type state: Dict
When the green agent requests to execute the browser application, and that request fails, this reward
component will keep track of that information. In that case, it doesn't matter whether the last webpage
had a 200 status code, because there has been an unsuccessful request since.
"""
if last_action_response.request == ["network", "node", self._node, "application", "WebBrowser", "execute"]:
self._last_request_failed = last_action_response.response.status != "success"
# if agent couldn't even get as far as sending the request (because for example the node was off), then
# apply a penalty
if self._last_request_failed:
return -1.0
# If the last request did actually go through, then check if the webpage also loaded
web_browser_state = access_from_nested_dict(state, self.location_in_state)
if web_browser_state is NOT_PRESENT_IN_STATE or "history" not in web_browser_state:
_LOGGER.info(
_LOGGER.debug(
"Web browser reward could not be calculated because the web browser history on node",
f"{self._node} was not reported in the simulation state. Returning 0.0",
)
@@ -242,15 +258,117 @@ class WebpageUnavailablePenalty(AbstractReward):
return cls(node_hostname=node_hostname)
class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
"""Penalises the agent when the green db clients fail to connect to the database."""
def __init__(self, node_hostname: str) -> None:
"""
Initialise the reward component.
:param node_hostname: Hostname of the node where the database client sits.
:type node_hostname: str
"""
self._node: str = node_hostname
self.location_in_state: List[str] = ["network", "nodes", node_hostname, "applications", "DatabaseClient"]
self._last_request_failed: bool = False
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""
Calculate the reward based on current simulation state, and the recent agent action.
When the green agent requests to execute the database client application, and that request fails, this reward
component will keep track of that information. In that case, it doesn't matter whether the last successful
request returned was able to connect to the database server, because there has been an unsuccessful request
since.
"""
if last_action_response.request == ["network", "node", self._node, "application", "DatabaseClient", "execute"]:
self._last_request_failed = last_action_response.response.status != "success"
# if agent couldn't even get as far as sending the request (because for example the node was off), then
# apply a penalty
if self._last_request_failed:
return -1.0
# If the last request was actually sent, then check if the connection was established.
db_state = access_from_nested_dict(state, self.location_in_state)
if db_state is NOT_PRESENT_IN_STATE or "last_connection_successful" not in db_state:
_LOGGER.debug(f"Can't calculate reward for {self.__class__.__name__}")
last_connection_successful = db_state["last_connection_successful"]
if last_connection_successful is False:
return -1.0
elif last_connection_successful is True:
return 1.0
return 0.0
@classmethod
def from_config(cls, config: Dict) -> AbstractReward:
"""
Build the reward component object from config.
:param config: Configuration dictionary.
:type config: Dict
"""
node_hostname = config.get("node_hostname")
return cls(node_hostname=node_hostname)
class SharedReward(AbstractReward):
"""Adds another agent's reward to the overall reward."""
def __init__(self, agent_name: Optional[str] = None) -> None:
"""
Initialise the shared reward.
The agent_name is a placeholder value. It starts off as none, but it must be set before this reward can work
correctly.
:param agent_name: The name whose reward is an input
:type agent_name: Optional[str]
"""
self.agent_name = agent_name
"""Agent whose reward to track."""
def default_callback(agent_name: str) -> Never:
"""
Default callback to prevent calling this reward until it's properly initialised.
SharedReward should not be used until the game layer replaces self.callback with a reference to the
function that retrieves the desired agent's reward. Therefore, we define this default callback that raises
an error.
"""
raise RuntimeError("Attempted to calculate SharedReward but it was not initialised properly.")
self.callback: Callable[[str], float] = default_callback
"""Method that retrieves an agent's current reward given the agent's name."""
def calculate(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Simply access the other agent's reward and return it."""
return self.callback(self.agent_name)
@classmethod
def from_config(cls, config: Dict) -> "SharedReward":
"""
Build the SharedReward object from config.
:param config: Configuration dictionary
:type config: Dict
"""
agent_name = config.get("agent_name")
return cls(agent_name=agent_name)
class RewardFunction:
"""Manages the reward function for the agent."""
__rew_class_identifiers: Dict[str, Type[AbstractReward]] = {
rew_class_identifiers: Dict[str, Type[AbstractReward]] = {
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty,
"WEBPAGE_UNAVAILABLE_PENALTY": WebpageUnavailablePenalty,
"GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY": GreenAdminDatabaseUnreachablePenalty,
"SHARED_REWARD": SharedReward,
}
"""List of reward class identifiers."""
def __init__(self):
"""Initialise the reward function object."""
@@ -269,7 +387,7 @@ class RewardFunction:
"""
self.reward_components.append((component, weight))
def update(self, state: Dict) -> float:
def update(self, state: Dict, last_action_response: "AgentActionHistoryItem") -> float:
"""Calculate the overall reward for the current state.
:param state: The current state of the simulation.
@@ -279,7 +397,7 @@ class RewardFunction:
for comp_and_weight in self.reward_components:
comp = comp_and_weight[0]
weight = comp_and_weight[1]
total += weight * comp.calculate(state=state)
total += weight * comp.calculate(state=state, last_action_response=last_action_response)
self.current_reward = total
return self.current_reward
@@ -297,7 +415,7 @@ class RewardFunction:
for rew_component_cfg in config["reward_components"]:
rew_type = rew_component_cfg["type"]
weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type]
rew_class = cls.rew_class_identifiers[rew_type]
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}))
new.register_component(component=rew_instance, weight=weight)
return new

14
src/primaite/game/agent/scripted_agents.py
View File

@@ -1,14 +0,0 @@
"""Agents with predefined behaviours."""
from primaite.game.agent.interface import AbstractScriptedAgent
class GreenWebBrowsingAgent(AbstractScriptedAgent):
"""Scripted agent which attempts to send web requests to a target node."""
raise NotImplementedError
class RedDatabaseCorruptingAgent(AbstractScriptedAgent):
"""Scripted agent which attempts to corrupt the database of the target node."""
raise NotImplementedError

0
src/primaite/game/agent/scripted_agents/__init__.py Normal file
View File

55
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py Normal file
View File

@@ -0,0 +1,55 @@
import random
from typing import Dict, Tuple
from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent
class DataManipulationAgent(AbstractScriptedAgent):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
next_execution_timestep: int = 0
starting_node_idx: int = 0
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setup_agent()
def _set_next_execution_timestep(self, timestep: int) -> None:
"""Set the next execution timestep with a configured random variance.
:param timestep: The timestep to add variance to.
"""
random_timestep_increment = random.randint(
-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance
)
self.next_execution_timestep = timestep + random_timestep_increment
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
"""Waits until a specific timestep, then attempts to execute its data manipulation application.
:param obs: Current observation for this agent, not used in DataManipulationAgent
:type obs: ObsType
:param timestep: The current simulation timestep, used for scheduling actions
:type timestep: int
:return: Action formatted in CAOS format
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
return "DONOTHING", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0}
def setup_agent(self) -> None:
"""Set the next execution timestep when the episode resets."""
self._select_start_node()
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
def _select_start_node(self) -> None:
"""Set the starting starting node of the agent to be a random node from this agent's action manager."""
# we are assuming that every node in the node manager has a data manipulation application at idx 0
num_nodes = len(self.action_manager.node_names)
self.starting_node_idx = random.randint(0, num_nodes - 1)

87
src/primaite/game/agent/scripted_agents/probabilistic_agent.py Normal file
View File

@@ -0,0 +1,87 @@
"""Agents with predefined behaviours."""
from typing import Dict, Optional, Tuple
import numpy as np
import pydantic
from gymnasium.core import ObsType
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractScriptedAgent
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction
class ProbabilisticAgent(AbstractScriptedAgent):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
class Settings(pydantic.BaseModel):
"""Config schema for Probabilistic agent settings."""
model_config = pydantic.ConfigDict(extra="forbid")
"""Strict validation."""
action_probabilities: Dict[int, float]
"""Probability to perform each action in the action map. The sum of probabilities should sum to 1."""
random_seed: Optional[int] = None
"""Random seed. If set, each episode the agent will choose the same random sequence of actions."""
# TODO: give the option to still set a random seed, but have it vary each episode in a predictable way
# for example if the user sets seed 123, have it be 123 + episode_num, so that each ep it's the next seed.
@pydantic.field_validator("action_probabilities", mode="after")
@classmethod
def probabilities_sum_to_one(cls, v: Dict[int, float]) -> Dict[int, float]:
"""Make sure the probabilities sum to 1."""
if not abs(sum(v.values()) - 1) < 1e-6:
raise ValueError("Green action probabilities must sum to 1")
return v
@pydantic.field_validator("action_probabilities", mode="after")
@classmethod
def action_map_covered_correctly(cls, v: Dict[int, float]) -> Dict[int, float]:
"""Ensure that the keys of the probability dictionary cover all integers from 0 to N."""
if not all((i in v) for i in range(len(v))):
raise ValueError(
"Green action probabilities must be defined as a mapping where the keys are consecutive integers "
"from 0 to N."
)
return v
def __init__(
self,
agent_name: str,
action_space: Optional[ActionManager],
observation_space: Optional[ObservationManager],
reward_function: Optional[RewardFunction],
settings: Dict = {},
) -> None:
# If the action probabilities are not specified, create equal probabilities for all actions
if "action_probabilities" not in settings:
num_actions = len(action_space.action_map)
settings = {"action_probabilities": {i: 1 / num_actions for i in range(num_actions)}}
# If seed not specified, set it to None so that numpy chooses a random one.
settings.setdefault("random_seed")
self.settings = ProbabilisticAgent.Settings(**settings)
self.rng = np.random.default_rng(self.settings.random_seed)
# convert probabilities from
self.probabilities = np.asarray(list(self.settings.action_probabilities.values()))
super().__init__(agent_name, action_space, observation_space, reward_function)
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""
Sample the action space randomly.
The probability of each action is given by the corresponding index in ``self.probabilities``.
:param obs: Current observation for this agent, not used in ProbabilisticAgent
:type obs: ObsType
:param timestep: The current simulation timestep, not used in ProbabilisticAgent
:type timestep: int
:return: Action formatted in CAOS format
:rtype: Tuple[str, Dict]
"""
choice = self.rng.choice(len(self.action_manager.action_map), p=self.probabilities)
return self.action_manager.get_action(choice)

21
src/primaite/game/agent/scripted_agents/random_agent.py Normal file
View File

@@ -0,0 +1,21 @@
from typing import Dict, Tuple
from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent
class RandomAgent(AbstractScriptedAgent):
"""Agent that ignores its observation and acts completely at random."""
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""Sample the action space randomly.
:param obs: Current observation for this agent, not used in RandomAgent
:type obs: ObsType
:param timestep: The current simulation timestep, not used in RandomAgent
:type timestep: int
:return: Action formatted in CAOS format
:rtype: Tuple[str, Dict]
"""
return self.action_manager.get_action(self.action_manager.space.sample())

234
src/primaite/game/game.py
View File

@@ -1,22 +1,25 @@
"""PrimAITE game - Encapsulates the simulation and agents."""
from ipaddress import IPv4Address
from typing import Dict, List
from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict
from primaite import getLogger
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.data_manipulation_bot import DataManipulationAgent
from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent, RandomAgent
from primaite.game.agent.observations import ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.session.io import SessionIO, SessionIOSettings
from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction, SharedReward
from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
from primaite.game.science import graph_has_cycle, topological_sort
from primaite.simulator.network.hardware.base import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.nmne import set_nmne_config
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
@@ -40,6 +43,7 @@ APPLICATION_TYPES_MAPPING = {
"DataManipulationBot": DataManipulationBot,
"DoSBot": DoSBot,
}
"""List of available applications that can be installed on nodes in the PrimAITE Simulation."""
SERVICE_TYPES_MAPPING = {
"DNSClient": DNSClient,
@@ -51,6 +55,7 @@ SERVICE_TYPES_MAPPING = {
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
"""List of available services that can be installed on nodes in the PrimAITE Simulation."""
class PrimaiteGameOptions(BaseModel):
@@ -63,8 +68,13 @@ class PrimaiteGameOptions(BaseModel):
model_config = ConfigDict(extra="forbid")
max_episode_length: int = 256
"""Maximum number of episodes for the PrimAITE game."""
ports: List[str]
"""A whitelist of available ports in the simulation."""
protocols: List[str]
"""A whitelist of available protocols in the simulation."""
thresholds: Optional[Dict] = {}
"""A dict containing the thresholds used for determining what is acceptable during observations."""
class PrimaiteGame:
@@ -79,18 +89,15 @@ class PrimaiteGame:
self.simulation: Simulation = Simulation()
"""Simulation object with which the agents will interact."""
self.agents: List[AbstractAgent] = []
"""List of agents."""
self.agents: Dict[str, AbstractAgent] = {}
"""Mapping from agent name to agent object."""
self.rl_agents: List[ProxyAgent] = []
"""Subset of agent list including only the reinforcement learning agents."""
self.rl_agents: Dict[str, ProxyAgent] = {}
"""Subset of agents which are intended for reinforcement learning."""
self.step_counter: int = 0
"""Current timestep within the episode."""
self.episode_counter: int = 0
"""Current episode number."""
self.options: PrimaiteGameOptions
"""Special options that apply for the entire game."""
@@ -109,6 +116,9 @@ class PrimaiteGame:
self.save_step_metadata: bool = False
"""Whether to save the RL agents' action, environment state, and other data at every single step."""
self._reward_calculation_order: List[str] = [name for name in self.agents]
"""Agent order for reward evaluation, as some rewards can be dependent on other agents' rewards."""
def step(self):
"""
Perform one step of the simulation/agent loop.
@@ -129,40 +139,49 @@ class PrimaiteGame:
"""
_LOGGER.debug(f"Stepping. Step counter: {self.step_counter}")
# Get the current state of the simulation
sim_state = self.get_sim_state()
# Update agents' observations and rewards based on the current state
self.update_agents(sim_state)
if self.step_counter == 0:
state = self.get_sim_state()
for agent in self.agents.values():
agent.update_observation(state=state)
# Apply all actions to simulation as requests
agent_actions = self.apply_agent_actions() # noqa
self.apply_agent_actions()
# Advance timestep
self.advance_timestep()
# Get the current state of the simulation
sim_state = self.get_sim_state()
# Update agents' observations and rewards based on the current state, and the response from the last action
self.update_agents(state=sim_state)
def get_sim_state(self) -> Dict:
"""Get the current state of the simulation."""
return self.simulation.describe_state()
def update_agents(self, state: Dict) -> None:
"""Update agents' observations and rewards based on the current state."""
for agent in self.agents:
agent.update_observation(state)
agent.update_reward(state)
for agent_name in self._reward_calculation_order:
agent = self.agents[agent_name]
if self.step_counter > 0: # can't get reward before first action
agent.update_reward(state=state)
agent.update_observation(state=state) # order of this doesn't matter so just use reward order
agent.reward_function.total_reward += agent.reward_function.current_reward
def apply_agent_actions(self) -> None:
"""Apply all actions to simulation as requests."""
agent_actions = {}
for agent in self.agents:
for _, agent in self.agents.items():
obs = agent.observation_manager.current_observation
rew = agent.reward_function.current_reward
action_choice, options = agent.get_action(obs, rew)
agent_actions[agent.agent_name] = (action_choice, options)
request = agent.format_request(action_choice, options)
self.simulation.apply_request(request)
return agent_actions
action_choice, parameters = agent.get_action(obs, timestep=self.step_counter)
request = agent.format_request(action_choice, parameters)
response = self.simulation.apply_request(request)
agent.process_action_response(
timestep=self.step_counter,
action=action_choice,
parameters=parameters,
request=request,
response=response,
)
def advance_timestep(self) -> None:
"""Advance timestep."""
@@ -178,20 +197,14 @@ class PrimaiteGame:
return True
return False
def reset(self) -> None:
"""Reset the game, this will reset the simulation."""
self.episode_counter += 1
self.step_counter = 0
_LOGGER.debug(f"Resetting primaite game, episode = {self.episode_counter}")
self.simulation.reset_component_for_episode(episode=self.episode_counter)
for agent in self.agents:
agent.reward_function.total_reward = 0.0
agent.reset_agent_for_episode()
def close(self) -> None:
"""Close the game, this will close the simulation."""
return NotImplemented
def setup_for_episode(self, episode: int) -> None:
"""Perform any final configuration of components to make them ready for the game to start."""
self.simulation.setup_for_episode(episode=episode)
@classmethod
def from_config(cls, cfg: Dict) -> "PrimaiteGame":
"""Create a PrimaiteGame object from a config dictionary.
@@ -209,10 +222,6 @@ class PrimaiteGame:
:return: A PrimaiteGame object.
:rtype: PrimaiteGame
"""
io_settings = cfg.get("io_settings", {})
_ = SessionIO(SessionIOSettings(**io_settings))
# Instantiating this ensures that the game saves to the correct output dir even without being part of a session
game = cls()
game.options = PrimaiteGameOptions(**cfg["game"])
game.save_step_metadata = cfg.get("io_settings", {}).get("save_step_metadata") or False
@@ -221,8 +230,12 @@ class PrimaiteGame:
sim = game.simulation
net = sim.network
nodes_cfg = cfg["simulation"]["network"]["nodes"]
links_cfg = cfg["simulation"]["network"]["links"]
simulation_config = cfg.get("simulation", {})
network_config = simulation_config.get("network", {})
nodes_cfg = network_config.get("nodes", [])
links_cfg = network_config.get("links", [])
for node_cfg in nodes_cfg:
node_ref = node_cfg["ref"]
n_type = node_cfg["type"]
@@ -230,28 +243,36 @@ class PrimaiteGame:
new_node = Computer(
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=node_cfg["subnet_mask"],
subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")),
default_gateway=node_cfg["default_gateway"],
dns_server=node_cfg["dns_server"],
operating_state=NodeOperatingState.ON,
dns_server=node_cfg.get("dns_server", None),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "server":
new_node = Server(
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=node_cfg["subnet_mask"],
subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")),
default_gateway=node_cfg["default_gateway"],
dns_server=node_cfg.get("dns_server"),
operating_state=NodeOperatingState.ON,
dns_server=node_cfg.get("dns_server", None),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "switch":
new_node = Switch(
hostname=node_cfg["hostname"],
num_ports=node_cfg.get("num_ports"),
operating_state=NodeOperatingState.ON,
num_ports=int(node_cfg.get("num_ports", "8")),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "router":
new_node = Router.from_config(node_cfg)
elif n_type == "firewall":
new_node = Firewall.from_config(node_cfg)
else:
_LOGGER.warning(f"invalid node type {n_type} in config")
if "services" in node_cfg:
@@ -264,8 +285,13 @@ class PrimaiteGame:
new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service.uuid
# start the service
new_service.start()
else:
_LOGGER.warning(f"service type not found {service_type}")
msg = f"Configuration contains an invalid service type: {service_type}"
_LOGGER.error(msg)
raise ValueError(msg)
# service-dependent options
if service_type == "DNSClient":
if "options" in service_cfg:
@@ -281,18 +307,16 @@ class PrimaiteGame:
if service_type == "DatabaseService":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.password = opt.get("db_password", None)
new_service.configure_backup(backup_server=IPv4Address(opt.get("backup_server_ip")))
new_service.start()
if service_type == "FTPServer":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.server_password = opt.get("server_password")
new_service.start()
if service_type == "NTPClient":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.ntp_server = IPv4Address(opt.get("ntp_server_ip"))
new_service.start()
if "applications" in node_cfg:
for application_cfg in node_cfg["applications"]:
new_application = None
@@ -304,7 +328,12 @@ class PrimaiteGame:
new_application = new_node.software_manager.software[application_type]
game.ref_map_applications[application_ref] = new_application.uuid
else:
_LOGGER.warning(f"application type not found {application_type}")
msg = f"Configuration contains an invalid application type: {application_type}"
_LOGGER.error(msg)
raise ValueError(msg)
# run the application
new_application.run()
if application_type == "DataManipulationBot":
if "options" in application_cfg:
@@ -312,7 +341,7 @@ class PrimaiteGame:
new_application.configure(
server_ip_address=IPv4Address(opt.get("server_ip")),
server_password=opt.get("server_password"),
payload=opt.get("payload"),
payload=opt.get("payload", "DELETE"),
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
data_manipulation_p_of_success=float(opt.get("data_manipulation_p_of_success", "0.1")),
)
@@ -327,7 +356,6 @@ class PrimaiteGame:
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.target_url = opt.get("target_url")
elif application_type == "DoSBot":
if "options" in application_cfg:
opt = application_cfg["options"]
@@ -344,10 +372,20 @@ class PrimaiteGame:
for nic_num, nic_cfg in node_cfg["network_interfaces"].items():
new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"]))
# temporarily set to 0 so all nodes are initially on
new_node.start_up_duration = 0
new_node.shut_down_duration = 0
net.add_node(new_node)
new_node.power_on()
# run through the power on step if the node is to be turned on at the start
if new_node.operating_state == NodeOperatingState.ON:
new_node.power_on()
game.ref_map_nodes[node_ref] = new_node.uuid
# set start up and shut down duration
new_node.start_up_duration = int(node_cfg.get("start_up_duration", 3))
new_node.shut_down_duration = int(node_cfg.get("shut_down_duration", 3))
# 2. create links between nodes
for link_cfg in links_cfg:
node_a = net.nodes[game.ref_map_nodes[link_cfg["endpoint_a_ref"]]]
@@ -364,7 +402,7 @@ class PrimaiteGame:
game.ref_map_links[link_cfg["ref"]] = new_link.uuid
# 3. create agents
agents_cfg = cfg["agents"]
agents_cfg = cfg.get("agents", [])
for agent_cfg in agents_cfg:
agent_ref = agent_cfg["ref"] # noqa: F841
@@ -382,21 +420,19 @@ class PrimaiteGame:
# CREATE REWARD FUNCTION
reward_function = RewardFunction.from_config(reward_function_cfg)
# OTHER AGENT SETTINGS
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
# CREATE AGENT
if agent_type == "GreenWebBrowsingAgent":
if agent_type == "ProbabilisticAgent":
# TODO: implement non-random agents and fix this parsing
new_agent = RandomAgent(
settings = agent_cfg.get("agent_settings", {})
new_agent = ProbabilisticAgent(
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=reward_function,
agent_settings=agent_settings,
settings=settings,
)
game.agents.append(new_agent)
elif agent_type == "ProxyAgent":
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
new_agent = ProxyAgent(
agent_name=agent_cfg["ref"],
action_space=action_space,
@@ -404,9 +440,10 @@ class PrimaiteGame:
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
game.rl_agents.append(new_agent)
game.rl_agents[agent_cfg["ref"]] = new_agent
elif agent_type == "RedDatabaseCorruptingAgent":
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
new_agent = DataManipulationAgent(
agent_name=agent_cfg["ref"],
action_space=action_space,
@@ -414,10 +451,55 @@ class PrimaiteGame:
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
else:
_LOGGER.warning(f"agent type {agent_type} not found")
msg = f"Configuration error: {agent_type} is not a valid agent type."
_LOGGER.error(msg)
raise ValueError(msg)
game.agents[agent_cfg["ref"]] = new_agent
game.simulation.set_original_state()
# Validate that if any agents are sharing rewards, they aren't forming an infinite loop.
game.setup_reward_sharing()
# Set the NMNE capture config
set_nmne_config(network_config.get("nmne_config", {}))
game.update_agents(game.get_sim_state())
return game
def setup_reward_sharing(self):
"""Do necessary setup to enable reward sharing between agents.
This method ensures that there are no cycles in the reward sharing. A cycle would be for example if agent_1
depends on agent_2 and agent_2 depends on agent_1. It would cause an infinite loop.
Also, SharedReward requires us to pass it a callback method that will provide the reward of the agent who is
sharing their reward. This callback is provided by this setup method.
Finally, this method sorts the agents in order in which rewards will be evaluated to make sure that any rewards
that rely on the value of another reward are evaluated later.
:raises RuntimeError: If the reward sharing is specified with a cyclic dependency.
"""
# construct dependency graph in the reward sharing between agents.
graph = {}
for name, agent in self.agents.items():
graph[name] = set()
for comp, weight in agent.reward_function.reward_components:
if isinstance(comp, SharedReward):
comp: SharedReward
graph[name].add(comp.agent_name)
# while constructing the graph, we might as well set up the reward sharing itself.
comp.callback = lambda agent_name: self.agents[agent_name].reward_function.current_reward
# make sure the graph is acyclic. Otherwise we will enter an infinite loop of reward sharing.
if graph_has_cycle(graph):
raise RuntimeError(
(
"Detected cycle in agent reward sharing. Check the agent reward function ",
"configuration: reward sharing can only go one way.",
)
)
# sort the agents so the rewards that depend on other rewards are always evaluated later
self._reward_calculation_order = topological_sort(graph)

78
src/primaite/game/science.py
View File

@@ -1,4 +1,5 @@
from random import random
from typing import Any, Iterable, Mapping
def simulate_trial(p_of_success: float) -> bool:
@@ -14,3 +15,80 @@ def simulate_trial(p_of_success: float) -> bool:
:returns: True if the trial is successful (with probability 'p_of_success'); otherwise, False.
"""
return random() < p_of_success
def graph_has_cycle(graph: Mapping[Any, Iterable[Any]]) -> bool:
"""Detect cycles in a directed graph.
Provide the graph as a dictionary that describes which nodes are linked. For example:
{0: {1,2}, 1:{2,3}, 3:{0}} here there's a cycle 0 -> 1 -> 3 -> 0
{'a': ('b','c'), c:('b')} here there is no cycle
:param graph: a mapping from node to a set of nodes to which it is connected.
:type graph: Mapping[Any, Iterable[Any]]
:return: Whether the graph has any cycles
:rtype: bool
"""
visited = set()
currently_visiting = set()
def depth_first_search(node: Any) -> bool:
"""Perform depth-first search (DFS) traversal to detect cycles starting from a given node."""
if node in currently_visiting:
return True # Cycle detected
if node in visited:
return False # Already visited, no need to explore further
visited.add(node)
currently_visiting.add(node)
for neighbour in graph.get(node, []):
if depth_first_search(neighbour):
return True # Cycle detected
currently_visiting.remove(node)
return False
# Start DFS traversal from each node
for node in graph:
if depth_first_search(node):
return True # Cycle detected
return False # No cycles found
def topological_sort(graph: Mapping[Any, Iterable[Any]]) -> Iterable[Any]:
"""
Perform topological sorting on a directed graph.
This guarantees that if there's a directed edge from node A to node B, then A appears before B.
:param graph: A dictionary representing the directed graph, where keys are node identifiers
and values are lists of outgoing edges from each node.
:type graph: dict[int, list[Any]]
:return: A topologically sorted list of node identifiers.
:rtype: list[Any]
"""
visited: set[Any] = set()
stack: list[Any] = []
def dfs(node: Any) -> None:
"""
Depth-first search traversal to visit nodes and their neighbors.
:param node: The current node to visit.
:type node: Any
"""
if node in visited:
return
visited.add(node)
for neighbour in graph.get(node, []):
dfs(neighbour)
stack.append(node)
# Perform DFS traversal from each node
for node in graph:
dfs(node)
return stack

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

46
src/primaite/interface/request.py Normal file
View File

@@ -0,0 +1,46 @@
from typing import Dict, ForwardRef, List, Literal, Union
from pydantic import BaseModel, ConfigDict, StrictBool, validate_call
RequestFormat = List[Union[str, int, float]]
RequestResponse = ForwardRef("RequestResponse")
"""This makes it possible to type-hint RequestResponse.from_bool return type."""
class RequestResponse(BaseModel):
"""Schema for generic request responses."""
model_config = ConfigDict(extra="forbid", strict=True)
"""Cannot have extra fields in the response. Anything custom goes into the data field."""
status: Literal["pending", "success", "failure", "unreachable"] = "pending"
"""
What is the current status of the request:
- pending - the request has not been received yet, or it has been received but it's still being processed.
- success - the request has been received and executed successfully.
- failure - the request has been received and attempted, but execution failed.
- unreachable - the request could not reach it's intended target, either because it doesn't exist or the target
is off.
"""
data: Dict = {}
"""Catch-all place to provide any additional data that was generated as a response to the request."""
# TODO: currently, status and data have default values, because I don't want to interrupt existing functionality too
# much. However, in the future we might consider making them mandatory.
@classmethod
@validate_call
def from_bool(cls, status_bool: StrictBool) -> RequestResponse:
"""
Construct a basic request response from a boolean.
True maps to a success status. False maps to a failure status.
:param status_bool: Whether to create a successful response
:type status_bool: bool
"""
if status_bool is True:
return cls(status="success", data={})
elif status_bool is False:
return cls(status="failure", data={})

4
src/primaite/main.py
View File

@@ -5,7 +5,7 @@ from pathlib import Path
from typing import Optional, Union
from primaite import getLogger
from primaite.config.load import example_config_path, load
from primaite.config.load import data_manipulation_config_path, load
from primaite.session.session import PrimaiteSession
# from primaite.primaite_session import PrimaiteSession
@@ -42,6 +42,6 @@ if __name__ == "__main__":
args = parser.parse_args()
if not args.config:
args.config = example_config_path()
args.config = data_manipulation_config_path()
run(args.config)

447
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb Normal file
View File

@@ -0,0 +1,447 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Customising red agents\n",
"\n",
"This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
"\n",
"First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
"\n",
"*(For a full explanation of the Data Manipulation scenario, check out the notebook `Data-Manipulation-E2E-Demonstration.ipynb`)*"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Imports\n",
"\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"import yaml\n",
"from pprint import pprint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def make_cfg_have_flat_obs(cfg):\n",
" for agent in cfg['agents']:\n",
" if agent['type'] == \"ProxyAgent\":\n",
" agent['agent_settings']['flatten_obs'] = False"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" make_cfg_have_flat_obs(cfg)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"obs, info = env.reset()\n",
"print('env created successfully')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'NODE_APPLICATION_EXECUTE':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By default, the red agent can start on client 1 or client 2. It starts its attack on a random step between 20 and 30, and it repeats its attack every 15-25 steps.\n",
"\n",
"It also has a 20% chance to fail to perform the port scan, and a 20% chance to fail launching the SQL attack. However it will continue where it left off after a failed step. I.e. if lucky, it can perform the port scan and SQL attack on the first try. If the port scan works, but the sql attack fails the first time it tries to attack, the next time it will not need to port scan again, it can go straight to trying to use SQL attack again."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for step in range(35):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since the agent does nothing most of the time, let's only print the steps where it performs an attack."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"for step in range(100):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Red Configuration\n",
"\n",
"There are two important parts of the YAML config for varying red agent behaviour.\n",
"\n",
"### Red agent settings\n",
"Here is an annotated config for the red agent in the data manipulation scenario.\n",
"```yaml\n",
" - ref: data_manipulation_attacker # name of agent\n",
" team: RED # not used, just for human reference\n",
" type: RedDatabaseCorruptingAgent # type of agent - this lets primaite know which agent class to use\n",
"\n",
" # Since the agent does not need to react to what is happening in the environment, the observation space is empty.\n",
" observation_space:\n",
" type: UC2RedObservation\n",
" options:\n",
" nodes: {}\n",
"\n",
" action_space:\n",
"\n",
" # The agent has two action choices, either do nothing, or execute a pre-scripted attack by using \n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
"\n",
" # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
" options:\n",
" nodes:\n",
" - node_name: client_1 # The network should have a node called client_1\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_1 should have DataManipulationBot configured on it\n",
" - node_name: client_2 # The network should have a node called client_2\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_2 should have DataManipulationBot configured on it\n",
"\n",
" # not important\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\n",
" # red agent does not need a reward function\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
"\n",
" # These actions are passed to the RedDatabaseCorruptingAgent init method, they dictate the schedule of attacks\n",
" agent_settings:\n",
" start_settings:\n",
" start_step: 25 # first attack at step 25\n",
" frequency: 20 # attacks will happen every 20 steps (on average)\n",
" variance: 5 # the timing of attacks will vary by up to 5 steps earlier or later\n",
"```\n",
"\n",
"### Malicious application settings\n",
"The red agent uses an application called `DataManipulationBot` which leverages a node's `DatabaseClient` to send a malicious SQL query to the database server. Here's an annotated example of how this is configured in the yaml *(with impertinent config items omitted)*:\n",
"```yaml\n",
"simulation:\n",
" network:\n",
" nodes:\n",
" - ref: client_1\n",
" hostname: client_1\n",
" type: computer\n",
" ip_address: 192.168.10.21\n",
" subnet_mask: 255.255.255.0\n",
" default_gateway: 192.168.10.1\n",
" \n",
" # \n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" options:\n",
" port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
" data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
" payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
" server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient # Database client must be installed in order for DataManipulationBot to function\n",
" options:\n",
" db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Editing red agent settings\n",
"\n",
"### Removing randomness from attack timing\n",
"\n",
"We can make the attacks happen at completely predictable intervals if we edit the red agent's settings to set variance to 0."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"start_settings:\n",
" start_step: 25\n",
" frequency: 20\n",
" variance: 0\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n",
" agent['agent_settings'] = change\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for step in range(100):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Making the start node always the same\n",
"\n",
"Normally, the agent randomly chooses between the nodes in its action space to send attacks from:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Open the config without changing anything\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(12):\n",
" env.reset()\n",
" for step in range(31):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can make the agent always start on a node of our choice letting that be the only node in the agent's action space."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"action_space:\n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
" options:\n",
" nodes:\n",
" - node_name: client_1\n",
" applications:\n",
" - application_name: DataManipulationBot\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n",
" agent.update(change)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(12):\n",
" env.reset()\n",
" for step in range(31):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Make the attack less likely to succeed.\n",
"\n",
"We can change the success probabilities within the data manipulation bot application. When the attack succeeds, the reward goes down.\n",
"\n",
"Setting the probabilities to 1.0 means the attack always succeeds - the reward will always drop\n",
"\n",
"Setting the probabilities to 0.0 means the attack always fails - the reward will never drop."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Make attack always succeed.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" options:\n",
" port_scan_p_of_success: 1.0\n",
" data_manipulation_p_of_success: 1.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(5):\n",
" env.reset()\n",
" for step in range(36):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if step_num == 35:\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Make attack always fail.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" options:\n",
" port_scan_p_of_success: 0.0\n",
" data_manipulation_p_of_success: 0.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(5):\n",
" env.reset()\n",
" for step in range(36):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if step_num == 35:\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

258
src/primaite/notebooks/uc2_demo.ipynb → src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -13,7 +13,7 @@
"source": [
"## Scenario\n",
"\n",
"The network consists of an office subnet and a server subnet. Clients in the office access a website which fetches data from a database.\n",
"The network consists of an office subnet and a server subnet. Clients in the office access a website which fetches data from a database. Occasionally, admins need to access the database directly from the clients.\n",
"\n",
"[<img src=\"_package_data/uc2_network.png\" width=\"500\"/>](_package_data/uc2_network.png)\n",
"\n",
@@ -46,7 +46,9 @@
"source": [
"## Green agent\n",
"\n",
"There are green agents logged onto client 1 and client 2. They use the web browser to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available."
"There are green agents logged onto client 1 and client 2. They use the web browser to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available.\n",
"\n",
"Sometimes, the green agents send a request directly to the database to check that it is reachable."
]
},
{
@@ -55,7 +57,7 @@
"source": [
"## Red agent\n",
"\n",
"The red agent waits a bit then sends a DELETE query to the database from client 1. If the delete is successful, the database file is flagged as compromised to signal that data is not available.\n",
"At the start of every episode, the red agent randomly chooses either client 1 or client 2 to login to. It waits a bit then sends a DELETE query to the database from its chosen client. If the delete is successful, the database file is flagged as compromised to signal that data is not available.\n",
"\n",
"[<img src=\"_package_data/uc2_attack.png\" width=\"500\"/>](_package_data/uc2_attack.png)\n",
"\n",
@@ -68,7 +70,9 @@
"source": [
"## Blue agent\n",
"\n",
"The blue agent can view the entire network, but the health statuses of components are not updated until a scan is performed. The blue agent should restore the database file from backup after it was compromised. It can also prevent further attacks by blocking client 1 from sending the malicious SQL query to the database server. This can be done by implementing an ACL rule on the router."
"The blue agent can view the entire network, but the health statuses of components are not updated until a scan is performed. The blue agent should restore the database file from backup after it was compromised. It can also prevent further attacks by blocking the red agent client from sending the malicious SQL query to the database server. This can be done by implementing an ACL rule on the router.\n",
"\n",
"However, these rules will also impact greens' ability to check the database connection. The blue agent should only block the infected client, it should let the other client connect freely. Once the attack has begun, automated traffic monitoring will detect it as suspicious network traffic. The blue agent's observation space will show this as an increase in the number of malicious network events (NMNE) on one of the network interfaces. To achieve optimal reward, the agent should only block the client which has the non-zero outbound NMNE."
]
},
{
@@ -84,7 +88,7 @@
"source": [
"## Scripted agents:\n",
"### Red\n",
"The red agent sits on client 1 and uses an application called DataManipulationBot whose sole purpose is to send a DELETE query to the database.\n",
"The red agent sits on a client and uses an application called DataManipulationBot whose sole purpose is to send a DELETE query to the database.\n",
"The red agent can choose one of two action each timestep:\n",
"1. do nothing\n",
"2. execute the data manipulation application\n",
@@ -92,6 +96,7 @@
"- start time\n",
"- frequency\n",
"- variance\n",
"\n",
"Attacks start at a random timestep between (start_time - variance) and (start_time + variance). After each attack, another is attempted after a random delay between (frequency - variance) and (frequency + variance) timesteps.\n",
"\n",
"The data manipulation app itself has an element of randomness because the attack has a probability of success. The default is 0.8 to succeed with the port scan step and 0.8 to succeed with the attack itself.\n",
@@ -100,9 +105,11 @@
"The red agent does not use information about the state of the network to decide its action.\n",
"\n",
"### Green\n",
"The green agents use the web browser application to send requests to the web server. The schedule of each green agent is currently random, meaning it will request webpage with a 50% probability, and do nothing with a 50% probability.\n",
"The green agents use the web browser application to send requests to the web server. The schedule of each green agent is currently random, it will do nothing 30% of the time, send a web request 60% of the time, and send a db status check 10% of the time.\n",
"\n",
"When a green agent is blocked from accessing the data through the webpage, this incurs a negative reward to the RL defender."
"When a green agent is blocked from accessing the data through the webpage, this incurs a negative reward to the RL defender.\n",
"\n",
"Also, when the green agent is blocked from checking the database status, it causes a small negative reward."
]
},
{
@@ -129,6 +136,9 @@
" - NETWORK_INTERFACES\n",
" - <nic_id 1-2>\n",
" - nic_status\n",
" - nmne\n",
" - inbound\n",
" - outbound\n",
" - operating_status\n",
"- LINKS\n",
" - <link_id 1-10>\n",
@@ -219,6 +229,14 @@
"|1|ENABLED|\n",
"|2|DISABLED|\n",
"\n",
"NMNE (number of malicious network events) means, for inbound or outbound traffic, means:\n",
"|value|NMNEs|\n",
"|--|--|\n",
"|0|None|\n",
"|1|1 - 5|\n",
"|2|6 - 10|\n",
"|3|More than 10|\n",
"\n",
"Link load has the following meaning:\n",
"|load|percent utilisation|\n",
"|--|--|\n",
@@ -289,11 +307,17 @@
"- `1`: Scan the web service - this refreshes the health status in the observation space\n",
"- `9`: Scan the database file - this refreshes the health status of the database file\n",
"- `13`: Patch the database service - This triggers the database to restore data from the backup server\n",
"- `19`: Shut down client 1\n",
"- `22`: Block outgoing traffic from client 1\n",
"- `26`: Block TCP traffic from client 1 to the database node\n",
"- `28-37`: Remove ACL rules 1-10\n",
"- `42`: Disconnect client 1 from the network\n",
"- `39`: Shut down client 1\n",
"- `40`: Start up client 1\n",
"- `46`: Block outgoing traffic from client 1\n",
"- `47`: Block outgoing traffic from client 2\n",
"- `50`: Block TCP traffic from client 1 to the database node\n",
"- `51`: Block TCP traffic from client 2 to the database node\n",
"- `52-61`: Remove ACL rules 1-10\n",
"- `66`: Disconnect client 1 from the network\n",
"- `67`: Reconnect client 1 to the network\n",
"- `68`: Disconnect client 2 from the network\n",
"- `69`: Reconnect client 2 to the network\n",
"\n",
"The other actions will either have no effect or will negatively impact the network, so the blue agent should avoid taking them."
]
@@ -304,9 +328,10 @@
"source": [
"## Reward Function\n",
"\n",
"The blue agent's reward is calculated using two measures:\n",
"The blue agent's reward is calculated using these measures:\n",
"1. Whether the database file is in a good state (+1 for good, -1 for corrupted, 0 for any other state)\n",
"2. Whether each green agents' most recent webpage request was successful (+1 for a `200` return code, -1 for a `404` return code and 0 otherwise).\n",
"3. Whether each green agents' most recent DB status check was successful (+1 for a successful connection, -1 for no connection).\n",
"\n",
"The file status reward and the two green-agent-related rewards are averaged to get a total step reward.\n"
]
@@ -346,9 +371,9 @@
"outputs": [],
"source": [
"# Imports\n",
"from primaite.config.load import example_config_path\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.game.game import PrimaiteGame\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"import yaml\n",
"from pprint import pprint\n"
]
@@ -365,21 +390,21 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"# create the env\n",
"with open(example_config_path(), 'r') as f:\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" # set success probability to 1.0 to avoid rerunning cells.\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
"game = PrimaiteGame.from_config(cfg)\n",
"env = PrimaiteGymEnv(game = game)\n",
"# Don't flatten obs as we are not training an agent and we wish to see the dict-formatted observations\n",
"env.agent.flatten_obs = False\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
" # don't flatten observations so that we can see what is going on\n",
" cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"obs, info = env.reset()\n",
"print('env created successfully')\n",
"pprint(obs)"
@@ -389,35 +414,49 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The red agent will start attacking at some point between step 20 and 30. When this happens, the reward will go from 1.0 to 0.0, and to -1.0 when the green agent tries to access the webpage."
"The red agent will start attacking at some point between step 20 and 30. When this happens, the reward will drop immediately, then drop to -0.8 when green agents try to access the webpage."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"for step in range(32):\n",
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'NODE_APPLICATION_EXECUTE':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for step in range(35):\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}, Blue reward:{reward}\" )"
" print(f\"step: {env.game.step_counter}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now the reward is -1, let's have a look at blue agent's observation."
"Now the reward is -0.8, let's have a look at blue agent's observation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"pprint(obs['NODES'])"
@@ -433,9 +472,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(9) # scan database file\n",
@@ -451,6 +488,13 @@
"File 1 in folder 1 on node 3 has `health_status = 2`, indicating that the database file is compromised."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Also, the NMNE outbound of either client 1 (node 6) or client 2 (node 7) increased from 0 to 1, but only right after the red attack, so we probably cannot see it now."
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -461,16 +505,14 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'].action}\" )\n",
"print(f\"Blue reward:{reward}\" )"
]
},
@@ -480,65 +522,68 @@
"source": [
"The patching takes two steps, so the reward hasn't changed yet. Let's do nothing for another timestep, the reward should improve.\n",
"\n",
"The reward will be 0 as soon as the file finishes restoring. Then, the reward will increase to 1 when the green agent makes a request. (Because the webapp access part of the reward does not update until a successful request is made.)\n",
"The reward will increase slightly as soon as the file finishes restoring. Then, the reward will increase to 1 when both green agents make successful requests.\n",
"\n",
"Run the following cell until the green action is `NODE_APPLICATION_EXECUTE`, then the reward should become 1. If you run it enough times, another red attack will happen and the reward will drop again."
"Run the following cell until the green action is `NODE_APPLICATION_EXECUTE` for application 0, then the reward should become 1. If you run it enough times, another red attack will happen and the reward will drop again."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
"obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Blue reward:{reward}\" )"
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user']}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user']}\" )\n",
"print(f\"Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The blue agent can prevent attacks by implementing an ACL rule to stop client_1 from sending POSTGRES traffic to the database. (Let's also patch the database file to get the reward back up.)"
"The blue agent can prevent attacks by implementing an ACL rule to stop client_1 or client_2 from sending POSTGRES traffic to the database. (Let's also patch the database file to get the reward back up.)\n",
"\n",
"Let's block both clients from communicating directly with the database."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"env.step(13) # Patch the database\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}, Blue reward:{reward}\" )\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"env.step(26) # Block client 1\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}, Blue reward:{reward}\" )\n",
"env.step(50) # Block client 1\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"for step in range(30):\n",
"env.step(51) # Block client 2\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"while abs(reward - 0.8) > 1e-5:\n",
" obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}, Blue reward:{reward}\" )"
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
" if env.game.step_counter > 10000:\n",
" break # make sure there's no infinite loop if something went wrong"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, even though the red agent executes an attack, the reward stays at 1.0"
"Now, even though the red agent executes an attack, the reward will stay at 0.8."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's also have a look at the ACL observation to verify our new ACL rule at position 5."
"Let's also have a look at the ACL observation to verify our new ACL rule at positions 5 and 6."
]
},
{
@@ -549,11 +594,92 @@
"source": [
"obs['ACL']"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can slightly increase the reward by unblocking the client which isn't being used by the attacker. If node 6 has outbound NMNEs, let's unblock client 2, and if node 7 has outbound NMNEs, let's unblock client 1."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(58) # Remove the ACL rule that blocks client 1\n",
"env.step(57) # Remove the ACL rule that blocks client 2\n",
"\n",
"tries = 0\n",
"while True:\n",
" tries += 1\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
"\n",
" if obs['NODES'][6]['NICS'][1]['NMNE']['outbound'] == 1:\n",
" # client 1 has NMNEs, let's block it\n",
" obs, reward, terminated, truncated, info = env.step(50) # block client 1\n",
" print(\"blocking client 1\")\n",
" break\n",
" elif obs['NODES'][7]['NICS'][1]['NMNE']['outbound'] == 1:\n",
" # client 2 has NMNEs, so let's block it\n",
" obs, reward, terminated, truncated, info = env.step(51) # block client 2\n",
" print(\"blocking client 2\")\n",
" break\n",
" if tries>100:\n",
" print(\"Error: NMNE never increased\")\n",
" break\n",
"\n",
"env.step(13) # Patch the database\n",
"print()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, the reward will eventually increase to 0.9, even after red agent attempts subsequent attacks."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"for step in range(40):\n",
" obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Reset the environment, you can rerun the other cells to verify that the attack works the same every episode. (except the red agent will move between `client_1` and `client_2`.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": "venv",
"language": "python",
"name": "python3"
},
@@ -567,9 +693,9 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
"nbformat_minor": 2
}

11
src/primaite/notebooks/training_example_ray_multi_agent.ipynb → src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb
View File

@@ -35,7 +35,7 @@
"\n",
"# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n",
"# to copy the files to your user data path.\n",
"with open(PRIMAITE_PATHS.user_config_path / 'example_config/example_config_2_rl_agents.yaml', 'r') as f:\n",
"with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"ray.init(local_mode=True)"
@@ -60,7 +60,7 @@
" policies={'defender_1','defender_2'}, # These names are the same as the agents defined in the example config.\n",
" policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,\n",
" )\n",
" .environment(env=PrimaiteRayMARLEnv, env_config={\"cfg\":cfg})#, disable_env_checking=True)\n",
" .environment(env=PrimaiteRayMARLEnv, env_config=cfg)#, disable_env_checking=True)\n",
" .rollouts(num_rollout_workers=0)\n",
" .training(train_batch_size=128)\n",
" )\n"
@@ -88,6 +88,13 @@
" param_space=config\n",
").fit()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {

6
src/primaite/notebooks/training_example_ray_single_agent.ipynb → src/primaite/notebooks/Training-an-RLLib-Agent.ipynb
View File

@@ -16,7 +16,7 @@
"source": [
"from primaite.game.game import PrimaiteGame\n",
"import yaml\n",
"from primaite.config.load import example_config_path\n",
"from primaite.config.load import data_manipulation_config_path\n",
"\n",
"from primaite.session.environment import PrimaiteRayEnv\n",
"from ray.rllib.algorithms import ppo\n",
@@ -26,7 +26,7 @@
"\n",
"# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n",
"# to copy the files to your user data path.\n",
"with open(example_config_path(), 'r') as f:\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"ray.init(local_mode=True)\n"
@@ -54,7 +54,7 @@
"metadata": {},
"outputs": [],
"source": [
"env_config = {\"cfg\":cfg}\n",
"env_config = cfg\n",
"\n",
"config = (\n",
" PPOConfig()\n",

Some files were not shown because too many files have changed in this diff Show More