diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 2375a391..624c9ca4 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -14,31 +14,36 @@ parameters: - name: matrix type: object default: - # - job_name: 'UbuntuPython38' - # py: '3.8' - # img: 'ubuntu-latest' - # every_time: false - # publish_coverage: false - - job_name: 'UbuntuPython311' - py: '3.11' + - job_name: 'UbuntuPython39' + py: '3.9' + img: 'ubuntu-latest' + every_time: false + publish_coverage: false + - job_name: 'UbuntuPython310' + py: '3.10' img: 'ubuntu-latest' every_time: true publish_coverage: true - # - job_name: 'WindowsPython38' - # py: '3.8' - # img: 'windows-latest' - # every_time: false - # publish_coverage: false + - job_name: 'UbuntuPython311' + py: '3.11' + img: 'ubuntu-latest' + every_time: false + publish_coverage: false + - job_name: 'WindowsPython39' + py: '3.9' + img: 'windows-latest' + every_time: false + publish_coverage: false - job_name: 'WindowsPython311' py: '3.11' img: 'windows-latest' every_time: false publish_coverage: false - # - job_name: 'MacOSPython38' - # py: '3.8' - # img: 'macOS-latest' - # every_time: false - # publish_coverage: false + - job_name: 'MacOSPython39' + py: '3.9' + img: 'macOS-latest' + every_time: false + publish_coverage: false - job_name: 'MacOSPython311' py: '3.11' img: 'macOS-latest' @@ -63,7 +68,7 @@ stages: displayName: 'Use Python ${{ item.py }}' - script: | - python -m pip install pre-commit + python -m pip install pre-commit>=6.1 pre-commit install pre-commit run --all-files displayName: 'Run pre-commits' @@ -71,7 +76,6 @@ stages: - script: | python -m pip install --upgrade pip==23.0.1 pip install wheel==0.38.4 --upgrade - pip install setuptools==66 --upgrade pip install build==0.10.0 pip install pytest-azurepipelines displayName: 'Install build dependencies' diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3088dc1d..d004dd6c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,10 +1,10 @@ repos: - - repo: local - hooks: - - id: ensure-copyright-clause - name: ensure copyright clause - entry: python copyright_clause_pre_commit_hook.py - language: python + # - repo: local + # hooks: + # - id: ensure-copyright-clause + # name: ensure copyright clause + # entry: python copyright_clause_pre_commit_hook.py + # language: python - repo: http://github.com/pre-commit/pre-commit-hooks rev: v4.4.0 hooks: @@ -31,7 +31,7 @@ repos: - id: isort args: [ "--profile", "black" ] - repo: http://github.com/PyCQA/flake8 - rev: 6.0.0 + rev: 6.1.0 hooks: - id: flake8 additional_dependencies: diff --git a/CHANGELOG.md b/CHANGELOG.md index e9147947..315579d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,24 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [4.0.0] = TBC + +### Added + +### Changed +- Agents now follow a common configuration format, simplifying the configuration of agents and their extensibilty. +- Actions within PrimAITE are now extensible, allowing for plugin support. +- Added a config schema to `ObservationManager`, `ActionManager`, and `RewardFunction`. +- Streamlined the way agents are created from config +- Agent config no longer requires a dummy action space if the action space is empty, the same applies for observation space and reward function +- Actions now support a config schema, to allow yaml data validation and default parameter values +- Action parameters are no longer defined through IDs, instead meaningful data is provided directly in the action map +- Test and example YAMLs have been updated to match the new agent and action schemas, such as: + - Removed empty action spaces, observation spaces, or reward spaces for agent which didn't use them + - Relabeled action parameters to match the new action config schemas, and updated the values to no longer rely on indices + - Removed action space options which were previously used for assigning meaning to action space IDs +- Updated tests that don't use YAMLs to still use the new action and agent schemas + ## [3.3.0] - 2024-09-04 ### Added diff --git a/docs/index.rst b/docs/index.rst index 42cc1d6d..aa7d16e0 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -70,7 +70,7 @@ PrimAITE incorporates the following features: - Architected with a separate Simulation layer and Game layer. This separation of concerns defines a clear path towards transfer learning with environments of differing fidelity; - Ability to reconfigure an RL reward function based on (a) the ability to counter the modelled adversarial cyber-attack, and (b) the ability to ensure success for green agents; -- Access Control List (ACL) functions for network devices (routers and firewalls), following standard ACL rule format (e.g., DENY / ALLOW, source / destination IP addresses, protocol and port); +- Access Control List (ACL) functions for network devices (routers and firewalls), following standard ACL rule format (e.g., DENY / PERMIT, source / destination IP addresses, protocol and port); - Application of traffic to the links of the system laydown adheres to the ACL rulesets and routing tables contained within each network device; - Provides RL environments adherent to the Farama Foundation Gymnasium (Previously OpenAI Gym) API, allowing integration with any compliant RL Agent frameworks; - Provides RL environments adherent to Ray RLlib environment specifications for single-agent and multi-agent scenarios; diff --git a/docs/source/about.rst b/docs/source/about.rst index da87102a..839bbb0b 100644 --- a/docs/source/about.rst +++ b/docs/source/about.rst @@ -184,7 +184,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! - 192.168.1.5 - ANY - ANY - All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or ALLOW for the same parameters will override an earlier entry. + All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or PERMIT for the same parameters will override an earlier entry. Observation Spaces ****************** The observation space provides the blue agent with information about the current status of nodes and links. @@ -331,7 +331,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! * Dictionary item {... ,1: [x1, x2, x3, x4, x5, x6] ...} The placeholders inside the list under the key '1' mean the following: * [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - * [0, 1] - Permission (0 = DENY, 1 = ALLOW) + * [0, 1] - Permission (0 = DENY, 1 = PERMIT) * [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) * [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) * [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) diff --git a/docs/source/action_masking.rst b/docs/source/action_masking.rst index dad6a484..4331b090 100644 --- a/docs/source/action_masking.rst +++ b/docs/source/action_masking.rst @@ -23,123 +23,117 @@ The following logic is applied: +------------------------------------------+---------------------------------------------------------------------+ | Action | Action Mask Logic | +==========================================+=====================================================================+ -| **DONOTHING** | Always Possible. | +| **do_nothing** | Always Possible. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_SCAN** | Node is on. Service is running. | +| **node_service_scan** | Node is on. Service is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_STOP** | Node is on. Service is running. | +| **node_service_stop** | Node is on. Service is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_START** | Node is on. Service is stopped. | +| **node_service_start** | Node is on. Service is stopped. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_PAUSE** | Node is on. Service is running. | +| **node_service_pause** | Node is on. Service is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_RESUME** | Node is on. Service is paused. | +| **node_service_resume** | Node is on. Service is paused. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_RESTART** | Node is on. Service is running. | +| **node_service_restart** | Node is on. Service is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_DISABLE** | Node is on. | +| **node_service_disable** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_ENABLE** | Node is on. Service is disabled. | +| **node_service_enable** | Node is on. Service is disabled. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SERVICE_FIX** | Node is on. Service is running. | +| **node_service_fix** | Node is on. Service is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_EXECUTE** | Node is on. | +| **node_application_execute** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_SCAN** | Node is on. Application is running. | +| **node_application_scan** | Node is on. Application is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_CLOSE** | Node is on. Application is running. | +| **node_application_close** | Node is on. Application is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_FIX** | Node is on. Application is running. | +| **node_application_fix** | Node is on. Application is running. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_INSTALL** | Node is on. | +| **node_application_install** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_APPLICATION_REMOVE** | Node is on. | +| **node_application_remove** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_SCAN** | Node is on. File exists. File not deleted. | +| **node_file_scan** | Node is on. File exists. File not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_CREATE** | Node is on. | +| **node_file_create** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_CHECKHASH** | Node is on. File exists. File not deleted. | +| **node_file_checkhash** | Node is on. File exists. File not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_DELETE** | Node is on. File exists. | +| **node_file_delete** | Node is on. File exists. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_REPAIR** | Node is on. File exists. File not deleted. | +| **node_file_repair** | Node is on. File exists. File not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_RESTORE** | Node is on. File exists. File is deleted. | +| **node_file_restore** | Node is on. File exists. File is deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_CORRUPT** | Node is on. File exists. File not deleted. | +| **node_file_corrupt** | Node is on. File exists. File not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FILE_ACCESS** | Node is on. File exists. File not deleted. | +| **node_file_access** | Node is on. File exists. File not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FOLDER_CREATE** | Node is on. | +| **node_folder_create** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FOLDER_SCAN** | Node is on. Folder exists. Folder not deleted. | +| **node_folder_scan** | Node is on. Folder exists. Folder not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FOLDER_CHECKHASH** | Node is on. Folder exists. Folder not deleted. | +| **node_folder_checkhash** | Node is on. Folder exists. Folder not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FOLDER_REPAIR** | Node is on. Folder exists. Folder not deleted. | +| **node_folder_repair** | Node is on. Folder exists. Folder not deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_FOLDER_RESTORE** | Node is on. Folder exists. Folder is deleted. | +| **node_folder_restore** | Node is on. Folder exists. Folder is deleted. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_OS_SCAN** | Node is on. | +| **node_os_scan** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **HOST_NIC_ENABLE** | NIC is disabled. Node is on. | +| **host_nic_enable** | NIC is disabled. Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **HOST_NIC_DISABLE** | NIC is enabled. Node is on. | +| **host_nic_disable** | NIC is enabled. Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SHUTDOWN** | Node is on. | +| **node_shutdown** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_STARTUP** | Node is off. | +| **node_startup** | Node is off. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_RESET** | Node is on. | +| **node_reset** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_PING_SCAN** | Node is on. | +| **node_nmap_ping_scan** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_PORT_SCAN** | Node is on. | +| **node_nmap_port_scan** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_NETWORK_SERVICE_RECON** | Node is on. | +| **node_network_service_recon** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NETWORK_PORT_ENABLE** | Node is on. Router is on. | +| **network_port_enable** | Node is on. Router is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NETWORK_PORT_DISABLE** | Router is on. | +| **network_port_disable** | Router is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **ROUTER_ACL_ADDRULE** | Router is on. | +| **router_acl_addrule** | Router is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **ROUTER_ACL_REMOVERULE** | Router is on. | +| **router_acl_removerule** | Router is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **FIREWALL_ACL_ADDRULE** | Firewall is on. | +| **firewall_acl_addrule** | Firewall is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **FIREWALL_ACL_REMOVERULE** | Firewall is on. | +| **firewall_acl_removerule** | Firewall is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_PING_SCAN** | Node is on. | +| **configure_database_client** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_PORT_SCAN** | Node is on. | +| **configure_ransomware_script** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_NMAP_NETWORK_SERVICE_RECON** | Node is on. | +| **c2_server_ransomware_configure** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **CONFIGURE_DATABASE_CLIENT** | Node is on. | +| **configure_dos_bot** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **CONFIGURE_RANSOMWARE_SCRIPT** | Node is on. | +| **configure_c2_beacon** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **CONFIGURE_DOSBOT** | Node is on. | +| **c2_server_ransomware_launch** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **CONFIGURE_C2_BEACON** | Node is on. | +| **c2_server_terminal_command** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **C2_SERVER_RANSOMWARE_LAUNCH** | Node is on. | +| **c2_server_data_exfiltrate** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **C2_SERVER_RANSOMWARE_CONFIGURE** | Node is on. | +| **node_account_change_password** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **C2_SERVER_TERMINAL_COMMAND** | Node is on. | +| **node_session_remote_login** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **C2_SERVER_DATA_EXFILTRATE** | Node is on. | +| **node_session_remote_logoff** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_ACCOUNTS_CHANGE_PASSWORD** | Node is on. | -+------------------------------------------+---------------------------------------------------------------------+ -| **SSH_TO_REMOTE** | Node is on. | -+------------------------------------------+---------------------------------------------------------------------+ -| **SESSIONS_REMOTE_LOGOFF** | Node is on. | -+------------------------------------------+---------------------------------------------------------------------+ -| **NODE_SEND_REMOTE_COMMAND** | Node is on. | +| **node_send_remote_command** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ diff --git a/docs/source/configuration/agents.rst b/docs/source/configuration/agents.rst index d11f7892..cf2b618f 100644 --- a/docs/source/configuration/agents.rst +++ b/docs/source/configuration/agents.rst @@ -23,19 +23,6 @@ Agents can be scripted (deterministic and stochastic), or controlled by a reinfo 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 @@ -91,10 +78,6 @@ For more information see :py:mod:`primaite.game.agent.observations` 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`` ^^^^^^^^^^^^^^ diff --git a/docs/source/how_to_guides/extensible_actions.rst b/docs/source/how_to_guides/extensible_actions.rst new file mode 100644 index 00000000..93a6cf21 --- /dev/null +++ b/docs/source/how_to_guides/extensible_actions.rst @@ -0,0 +1,67 @@ +.. only:: comment + + © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + + +Extensible Actions +****************** + + +Changes to Actions class Structure. +=================================== + +Actions within PrimAITE have been updated to inherit from a base class, AbstractAction, standardising their format and allowing for easier creation of custom actions. Actions now use a ``ConfigSchema`` to define the possible configuration variables, and use pydantic to enforce correct parameters are passed through. + + +Developing Custom Actions. +========================== + +Custom actions within PrimAITE must be a sub-class of `AbstractAction`, and contain 3 key items: + +#. ConfigSchema class + +#. Unique Identifier + +#. `form_request` method. + + +ConfigSchema +############ + +The ConfigSchema sub-class of the action must contain all `configurable` variables within the action, that would be specified within the environments configuration YAML file. + + +Unique Identifier +################# + +When declaring a custom class, it must have a unique identifier string, that allows PrimAITE to generate the correct action when needed. + +.. code:: Python + + class CreateDirectoryAction(AbstractAction, identifier="node_folder_create") + + config: CreateDirectoryAction.ConfigSchema + + class ConfigSchema(AbstractAction.ConfigSchema): + + verb: ClassVar[str] = "create" + node_name: str + directory_name: str + + def form_request(cls, config: ConfigSchema) -> RequestFormat: + return ["network", + "node", + config.node_name, + "file_system", + config.verb, + "folder", + config.directory_name, + ] + +The above action would fail pydantic validation as the identifier "node_folder_create" is already used by the `NodeFolderCreateAction`, and would create a duplicate listing within `AbstractAction._registry`. + + +form_request method +################### + +PrimAITE actions need to have a `form_request` method, which can be passed to the `RequestManager` for processing. This allows the custom action to be actioned within the simulation environment. diff --git a/docs/source/how_to_guides/extensible_agents.rst b/docs/source/how_to_guides/extensible_agents.rst new file mode 100644 index 00000000..256e96ca --- /dev/null +++ b/docs/source/how_to_guides/extensible_agents.rst @@ -0,0 +1,78 @@ +.. only:: comment + + © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + +.. _about: + +Extensible Agents +***************** + +Agents defined within PrimAITE have been updated to allow for easier creation of new bespoke agents for use in custom environments. + + +Developing Agents for PrimAITE +============================== + +All agent types within PrimAITE must be subclassed from ``AbstractAgent`` in order to be used from configuration YAML files. This then allows you to implement any custom agent logic for the new agent in your training scenario. Examples of implementing custom agent logic can be seen in pre-existing agents, such as the ``DataManipulationBot`` and ``RandomAgent``. + +The core features that should be implemented in any new agent are detailed below: + +#. **ConfigSchema**: + + Configurable items within a new agent within PrimAITE should contain a ``ConfigSchema`` which holds all configurable variables of the agent. This should not include parameters related to its *state*, these would be listed seperately. + Agent generation will fail pydantic checks if incorrect or invalid parameters are passed to the ConfigSchema of the chosen Agent. + + + .. code-block:: python + + class ExampleAgent(AbstractAgent, identifier = "ExampleAgent"): + """An example agent for demonstration purposes.""" + + config: "ExampleAgent.ConfigSchema" = Field(default_factory= lambda: ExampleAgent.ConfigSchema()) + """Agent configuration""" + num_executions: int = 0 + """Number of action executions by agent""" + + class ConfigSchema(AbstractAgent.ConfigSchema): + """ExampleAgent configuration schema""" + + type: str = "ExampleAgent + """Name of agent""" + starting_host: int + """Host node that this agent should start from in the given environment.""" + + + .. code-block:: yaml + + - ref: example_green_agent + team: GREEN + type: ExampleAgent + + action_space: + action_map: + 0: + action: do_nothing + options: {} + reward_function: + reward_components: + - type: DUMMY + + agent_settings: + start_step: 25 + frequency: 20 + variance: 5 + starting_host: "Server_1" + + +#. **Identifiers**: + + All agent classes should have an ``identifier`` attribute, a unique kebab-case string, for when they are added to the base ``AbstractAgent`` registry. This is then specified in your configuration YAML, and used by PrimAITE to generate the correct Agent. + +Changes to YAML file +==================== + +PrimAITE v4.0.0 introduces some breaking changes to how environment configuration yaml files are created. YAML files created for Primaite versions 3.3.0 should be compatible through a translation function, though it is encouraged that these are updated to reflect the updated format of 4.0.0+. + +Agents now follow a more standardised settings definition, so should be more consistent across YAML files and the available agent types with PrimAITE. + +All configurable items for agents sit under the ``agent_settings`` heading within your YAML files. There is no need for the inclusion of a ``start_settings``. Please see the above YAML example for full changes to agents. diff --git a/docs/source/primaite-dependencies.rst b/docs/source/primaite-dependencies.rst index 14a96349..ce2087ca 100644 --- a/docs/source/primaite-dependencies.rst +++ b/docs/source/primaite-dependencies.rst @@ -2,44 +2,44 @@ © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| Name | Version | License | Description | URL | -+===================+=========+====================================+=======================================================================================================+====================================================================+ -| gymnasium | 0.28.1 | MIT License | A standard API for reinforcement learning and a diverse set of reference environments (formerly Gym). | https://farama.org | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| ipywidgets | 8.1.5 | BSD License | Jupyter interactive widgets | http://jupyter.org | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| jupyterlab | 3.6.1 | BSD License | JupyterLab computational environment | https://jupyter.org | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| kaleido | 0.2.1 | MIT | Static image export for web-based visualization libraries with zero dependencies | https://github.com/plotly/Kaleido | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| matplotlib | 3.7.1 | Python Software Foundation License | Python plotting package | https://matplotlib.org | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| networkx | 3.1 | BSD License | Python package for creating and manipulating graphs and networks | https://networkx.org/ | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| numpy | 1.23.5 | BSD License | NumPy is the fundamental package for array computing with Python. | https://www.numpy.org | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| platformdirs | 3.5.1 | MIT License | A small Python package for determining appropriate platform-specific dirs, e.g. a "user data dir". | https://github.com/platformdirs/platformdirs | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| plotly | 5.15.0 | MIT License | An open-source, interactive data visualization library for Python | https://plotly.com/python/ | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| polars | 0.20.30 | MIT License | Blazingly fast DataFrame library | https://www.pola.rs/ | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| prettytable | 3.8.0 | BSD License (BSD (3 clause)) | A simple Python library for easily displaying tabular data in a visually appealing ASCII table format | https://github.com/jazzband/prettytable | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| pydantic | 2.7.0 | MIT License | Data validation using Python type hints | https://github.com/pydantic/pydantic | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| PyYAML | 6.0 | MIT License | YAML parser and emitter for Python | https://pyyaml.org/ | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| ray | 2.32.0 | Apache 2.0 | Ray provides a simple, universal API for building distributed applications. | https://github.com/ray-project/ray | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| stable-baselines3 | 2.1.0 | MIT | Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms. | https://github.com/DLR-RM/stable-baselines3 | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| tensorflow | 2.12.0 | Apache Software License | TensorFlow is an open source machine learning framework for everyone. | https://www.tensorflow.org/ | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| typer | 0.9.0 | MIT License | Typer, build great CLIs. Easy to code. Based on Python type hints. | https://github.com/tiangolo/typer | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| Deepdiff | 8.0.1 | MIT License | Deep difference of dictionaries, iterables, strings, and any other object objects. | https://github.com/seperman/deepdiff | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ -| sb3_contrib | 2.1.0 | MIT License | Contrib package for Stable-Baselines3 - Experimental reinforcement learning (RL) code (Action Masking)| https://github.com/Stable-Baselines-Team/stable-baselines3-contrib | -+-------------------+---------+------------------------------------+-------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------+ ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| Name | Supported Version | Built Version | License | Description | URL | ++===================+=====================+===============+======================================+========================================================================================================+=====================================================================+ +| gymnasium | 0.28.1 | 0.28.1 | MIT License | A standard API for reinforcement learning and a diverse set of reference environments (formerly Gym). | https://farama.org | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ipywidgets | ~=8.0 | 8.1.5 | BSD License | Jupyter interactive widgets | http://jupyter.org | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| jupyterlab | 3.6.1 | 3.6.1 | BSD License | JupyterLab computational environment | https://jupyter.org | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| kaleido | ==0.2.1 | 0.2.1 | MIT | Static image export for web-based visualization libraries with zero dependencies | https://github.com/plotly/Kaleido | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| matplotlib | >=3.7.1 | 3.7.1 | Python Software Foundation License | Python plotting package | https://matplotlib.org | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| networkx | 3.1 | 3.1 | BSD License | Python package for creating and manipulating graphs and networks | https://networkx.org/ | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| numpy | ~1.23 | 1.23.5 | BSD License | NumPy is the fundamental package for array computing with Python. | https://www.numpy.org | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| platformdirs | 3.5.1 | 3.5.1 | MIT License | A small Python package for determining appropriate platform-specific dirs, e.g. a "user data dir". | https://github.com/platformdirs/platformdirs | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| plotly | 5.15 | 5.15.0 | MIT License | An open-source, interactive data visualization library for Python | https://plotly.com/python/ | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| polars | 0.20.30 | 0.20.30 | MIT License | Blazingly fast DataFrame library | https://www.pola.rs/ | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| prettytable | 3.8.0 | 3.8.0 | BSD License (BSD (3 clause)) | A simple Python library for easily displaying tabular data in a visually appealing ASCII table format | https://github.com/jazzband/prettytable | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| pydantic | 2.7.0 | 2.7.0 | MIT License | Data validation using Python type hints | https://github.com/pydantic/pydantic | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| PyYAML | >=6.0 | 6.0 | MIT License | YAML parser and emitter for Python | https://pyyaml.org/ | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ray | >=2.20, <2.33 | 2.32.0 | Apache 2.0 | Ray provides a simple, universal API for building distributed applications. | https://github.com/ray-project/ray | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| stable-baselines3 | 2.1.0 | 2.1.0 | MIT | Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms. | https://github.com/DLR-RM/stable-baselines3 | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| tensorflow | ~=2.12 | 2.12.0 | Apache Software License | TensorFlow is an open source machine learning framework for everyone. | https://www.tensorflow.org/ | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| typer | >=0.9 | 0.9.0 | MIT License | Typer, build great CLIs. Easy to code. Based on Python type hints. | https://github.com/tiangolo/typer | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| Deepdiff | 8.0.1 | 8.0.1 | MIT License | Deep difference of dictionaries, iterables, strings, and any other object objects. | https://github.com/seperman/deepdiff | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| sb3_contrib | 2.1.0 | 2.1.0 | MIT License | Contrib package for Stable-Baselines3 - Experimental reinforcement learning (RL) code (Action Masking) | https://github.com/Stable-Baselines-Team/stable-baselines3-contrib | ++-------------------+---------------------+---------------+--------------------------------------+--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------+ diff --git a/docs/source/simulation_components/system/applications/data_manipulation_bot.rst b/docs/source/simulation_components/system/applications/data_manipulation_bot.rst index 91c33ede..49dc3baf 100644 --- a/docs/source/simulation_components/system/applications/data_manipulation_bot.rst +++ b/docs/source/simulation_components/system/applications/data_manipulation_bot.rst @@ -113,18 +113,6 @@ If not using the data manipulation bot manually, it needs to be used with a data folders: {} action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client_1 - applications: - - application_ref: data_manipulation_bot - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - reward_function: reward_components: - type: DUMMY diff --git a/docs/source/simulation_components/system/applications/ransomware_script.rst b/docs/source/simulation_components/system/applications/ransomware_script.rst index b79ca802..192618fc 100644 --- a/docs/source/simulation_components/system/applications/ransomware_script.rst +++ b/docs/source/simulation_components/system/applications/ransomware_script.rst @@ -70,7 +70,7 @@ Python Configuration ============= -The RansomwareScript inherits configuration options such as ``fix_duration`` from its parent class. However, for the ``RansomwareScript`` the most relevant option is ``server_ip``. +The RansomwareScript inherits configuration options such as ``fixing_duration`` from its parent class. However, for the ``RansomwareScript`` the most relevant option is ``server_ip``. ``server_ip`` diff --git a/docs/source/simulation_components/system/common/common_configuration.rst b/docs/source/simulation_components/system/common/common_configuration.rst index 411fd529..c1bbd4b2 100644 --- a/docs/source/simulation_components/system/common/common_configuration.rst +++ b/docs/source/simulation_components/system/common/common_configuration.rst @@ -22,8 +22,8 @@ options The configuration options are the attributes that fall under the options for an application or service. -fix_duration -"""""""""""" +fixing_duration +""""""""""""""" Optional. Default value is ``2``. diff --git a/pyproject.toml b/pyproject.toml index 354df8b2..e840797c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "primaite" description = "PrimAITE (Primary-level AI Training Environment) is a simulation environment for training AI under the ARCD programme." authors = [{name="Defence Science and Technology Laboratory UK", email="oss@dstl.gov.uk"}] license = {file = "LICENSE"} -requires-python = ">=3.9, <3.12" +requires-python = ">=3.9, <3.13" dynamic = ["version", "readme"] classifiers = [ "Development Status :: 5 - Production/Stable", @@ -26,15 +26,15 @@ dependencies = [ "gymnasium==0.28.1", "jupyterlab==3.6.1", "kaleido==0.2.1", - "matplotlib==3.7.1", + "matplotlib>=3.7.1", "networkx==3.1", - "numpy==1.23.5", + "numpy~=1.23", "platformdirs==3.5.1", "plotly==5.15.0", "polars==0.20.30", "prettytable==3.8.0", - "PyYAML==6.0", - "typer[all]==0.9.0", + "PyYAML>=6.0", + "typer[all]>=0.9", "pydantic==2.7.0", "ipywidgets", "deepdiff" @@ -53,8 +53,8 @@ license-files = ["LICENSE"] [project.optional-dependencies] rl = [ "ray[rllib] >= 2.20.0, <2.33", - "tensorflow==2.12.0", - "stable-baselines3[extra]==2.1.0", + "tensorflow~=2.12", + "stable-baselines3==2.1.0", "sb3-contrib==2.1.0", ] dev = [ @@ -69,7 +69,7 @@ dev = [ "pytest-xdist==3.3.1", "pytest-cov==4.0.0", "pytest-flake8==1.1.1", - "setuptools==66", + "setuptools==75.6.0", "Sphinx==7.1.2", "sphinx-copybutton==0.5.2", "wheel==0.38.4", diff --git a/run_test_and_coverage.py b/run_test_and_coverage.py index 3bd9072d..dfa71f74 100644 --- a/run_test_and_coverage.py +++ b/run_test_and_coverage.py @@ -1,4 +1,4 @@ -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK import subprocess import sys from typing import Any diff --git a/src/primaite/config/_package_data/data_manipulation.yaml b/src/primaite/config/_package_data/data_manipulation.yaml index 97442903..b0d5d087 100644 --- a/src/primaite/config/_package_data/data_manipulation.yaml +++ b/src/primaite/config/_package_data/data_manipulation.yaml @@ -30,35 +30,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -79,35 +66,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: WebBrowser reward_function: reward_components: @@ -128,33 +102,12 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - 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 + agent_settings: + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -208,8 +161,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -235,490 +188,426 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - reward_function: reward_components: diff --git a/src/primaite/config/_package_data/data_manipulation_marl.yaml b/src/primaite/config/_package_data/data_manipulation_marl.yaml index ba666781..e45f193e 100644 --- a/src/primaite/config/_package_data/data_manipulation_marl.yaml +++ b/src/primaite/config/_package_data/data_manipulation_marl.yaml @@ -26,35 +26,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -75,35 +62,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: WebBrowser reward_function: reward_components: @@ -124,36 +98,13 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - - type: NODE_FILE_DELETE - - type: NODE_FILE_CORRUPT - - type: NODE_OS_SCAN - 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 + agent_settings: + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender_1 team: BLUE @@ -202,8 +153,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -229,488 +180,423 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 + node_name: client_2 + nic_num: 1 reward_function: @@ -782,8 +668,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -809,494 +695,426 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - options: - target_router: router_1 - - type: ROUTER_ACL_REMOVERULE - options: - target_router: router_1 - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - reward_function: reward_components: - type: DATABASE_FILE_INTEGRITY diff --git a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml index b4457a28..e18de191 100644 --- a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml @@ -6,68 +6,48 @@ game: agents: - ref: RL_Agent type: ProxyAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_SHUTDOWN - - type: NODE_STARTUP - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 0 + node_name: client_1 2: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 1 + node_name: server 3: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: client_1 4: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: server 5: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 0 - nic_id: 0 + node_name: client_1 + nic_num: 1 6: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 1 - nic_id: 0 + node_name: server + nic_num: 1 7: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 0 - nic_id: 0 + node_name: client_1 + nic_num: 1 8: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 1 - nic_id: 0 - options: - nodes: - - node_name: client_1 - - node_name: server - max_folders_per_node: 0 - max_files_per_folder: 0 - max_services_per_node: 0 - max_nics_per_node: 1 - max_acl_rules: 0 - ip_list: - - 192.168.1.2 - - 192.168.1.3 - reward_function: - reward_components: [] + node_name: server + nic_num: 1 simulation: network: diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml index 98d2392a..677cd5a5 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml @@ -6,25 +6,17 @@ agents: &greens action_probabilities: 0: 0.2 1: 0.8 - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DatabaseClient action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client + application_name: DatabaseClient reward_function: reward_components: diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml index 17a5977b..eb7823f8 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml @@ -6,25 +6,17 @@ agents: &greens action_probabilities: 0: 0.95 1: 0.05 - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DatabaseClient action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client + application_name: DatabaseClient reward_function: reward_components: diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml index 31675a0b..0170143f 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml @@ -3,24 +3,9 @@ reds: &reds team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DataManipulationBot - - reward_function: - reward_components: - - type: DUMMY - agent_settings: - start_settings: - start_step: 10 - frequency: 10 - variance: 0 + possible_start_nodes: [client,] + target_application: DataManipulationBot + start_step: 10 + frequency: 10 + variance: 0 diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml index c5572b89..55bee3fb 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml @@ -3,24 +3,9 @@ reds: &reds team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DataManipulationBot - - reward_function: - reward_components: - - type: DUMMY - agent_settings: - start_settings: - start_step: 3 - frequency: 2 - variance: 1 + possible_start_nodes: [client_1] + target_application: DataManipulationBot + start_step: 3 + frequency: 2 + variance: 1 diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml index dfd200f3..c692c725 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml @@ -54,65 +54,46 @@ agents: - server:eth-1<->switch_1:eth-2 action_space: - action_list: - - type: DONOTHING - - type: NODE_SHUTDOWN - - type: NODE_STARTUP - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 0 + node_name: client_1 2: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 1 + node_name: server 3: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: client_1 4: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: server 5: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 0 - nic_id: 0 + node_name: client_1 + nic_num: 1 6: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 1 - nic_id: 0 + node_name: server + nic_num: 1 7: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 0 - nic_id: 0 + node_name: client_1 + nic_num: 1 8: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 1 - nic_id: 0 - options: - nodes: - - node_name: client - - node_name: server - - max_folders_per_node: 0 - max_files_per_folder: 0 - max_services_per_node: 0 - max_nics_per_node: 1 - max_acl_rules: 0 - ip_list: - - 192.168.1.2 - - 192.168.1.3 + node_name: server + nic_num: 1 reward_function: reward_components: diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py deleted file mode 100644 index 5ec122ec..00000000 --- a/src/primaite/game/agent/actions.py +++ /dev/null @@ -1,1804 +0,0 @@ -# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -""" -This module contains the ActionManager class which belongs to the Agent class. - -An agent's action space is made up of a collection of actions. Each action is an instance of a subclass of -AbstractAction. The ActionManager is responsible for: - 1. Creating the action space from a list of action types. - 2. Converting an integer action choice into a specific action and parameter choice. - 3. Converting an action and parameter choice into a request which can be ingested by the PrimAITE simulation. This - ensures that requests conform to the simulator's request format. -""" -import itertools -from abc import ABC, abstractmethod -from typing import Dict, List, Literal, Optional, Tuple, TYPE_CHECKING, Union - -from gymnasium import spaces -from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo - -from primaite import getLogger -from primaite.interface.request import RequestFormat - -_LOGGER = getLogger(__name__) - -if TYPE_CHECKING: - from primaite.game.game import PrimaiteGame - - -class AbstractAction(ABC): - """Base class for actions.""" - - @abstractmethod - def __init__(self, manager: "ActionManager", **kwargs) -> None: - """ - Init method for action. - - All action init functions should accept **kwargs as a way of ignoring extra arguments. - - Since many parameters are defined for the action space as a whole (such as max files per folder, max services - per node), we need to pass those options to every action that gets created. To prevent verbosity, these - parameters are just broadcasted to all actions and the actions can pay attention to the ones that apply. - """ - self.name: str = "" - """Human-readable action identifier used for printing, logging, and reporting.""" - self.shape: Dict[str, int] = {} - """Dictionary describing the number of options for each parameter of this action. The keys of this dict must - align with the keyword args of the form_request method.""" - self.manager: ActionManager = manager - """Reference to the ActionManager which created this action. This is used to access the game and simulation - objects.""" - - @abstractmethod - def form_request(self) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [] - - -class DoNothingAction(AbstractAction): - """Action which does nothing. This is here to allow agents to be idle if they choose to.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - self.name = "DONOTHING" - self.shape: Dict[str, int] = { - "dummy": 1, - } - # This action does not accept any parameters, therefore it technically has a gymnasium shape of Discrete(1), - # i.e. a choice between one option. To make enumerating this action easier, we are adding a 'dummy' paramter - # with one option. This just aids the Action Manager to enumerate all possibilities. - - def form_request(self, **kwargs) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return ["do_nothing"] - - -class NodeServiceAbstractAction(AbstractAction): - """ - Base class for service actions. - - Any action which applies to a service and uses node_id and service_id as its only two parameters can inherit from - this base class. - """ - - @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "service_id": num_services} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int, service_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - service_name = self.manager.get_service_name_by_idx(node_id, service_id) - if node_name is None or service_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "service", service_name, self.verb] - - -class NodeServiceScanAction(NodeServiceAbstractAction): - """Action which scans a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "scan" - - -class NodeServiceStopAction(NodeServiceAbstractAction): - """Action which stops a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "stop" - - -class NodeServiceStartAction(NodeServiceAbstractAction): - """Action which starts a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "start" - - -class NodeServicePauseAction(NodeServiceAbstractAction): - """Action which pauses a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "pause" - - -class NodeServiceResumeAction(NodeServiceAbstractAction): - """Action which resumes a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "resume" - - -class NodeServiceRestartAction(NodeServiceAbstractAction): - """Action which restarts a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "restart" - - -class NodeServiceDisableAction(NodeServiceAbstractAction): - """Action which disables a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "disable" - - -class NodeServiceEnableAction(NodeServiceAbstractAction): - """Action which enables a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "enable" - - -class NodeServiceFixAction(NodeServiceAbstractAction): - """Action which fixes a service.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) - self.verb: str = "fix" - - -class NodeApplicationAbstractAction(AbstractAction): - """ - Base class for application actions. - - Any action which applies to an application and uses node_id and application_id as its only two parameters can - inherit from this base class. - """ - - @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "application_id": num_applications} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int, application_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - application_name = self.manager.get_application_name_by_idx(node_id, application_id) - if node_name is None or application_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "application", application_name, self.verb] - - -class NodeApplicationExecuteAction(NodeApplicationAbstractAction): - """Action which executes an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications) - self.verb: str = "execute" - - -class NodeApplicationScanAction(NodeApplicationAbstractAction): - """Action which scans an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications) - self.verb: str = "scan" - - -class NodeApplicationCloseAction(NodeApplicationAbstractAction): - """Action which closes an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications) - self.verb: str = "close" - - -class NodeApplicationFixAction(NodeApplicationAbstractAction): - """Action which fixes an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications) - self.verb: str = "fix" - - -class NodeApplicationInstallAction(AbstractAction): - """Action which installs an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes} - - def form_request(self, node_id: int, application_name: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - return [ - "network", - "node", - node_name, - "software_manager", - "application", - "install", - application_name, - ] - - -class ConfigureDatabaseClientAction(AbstractAction): - """Action which sets config parameters for a database client on a node.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this action.""" - - model_config = ConfigDict(extra="forbid") - server_ip_address: Optional[str] = None - server_password: Optional[str] = None - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, config: Dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - ConfigureDatabaseClientAction._Opts.model_validate(config) # check that options adhere to schema - return ["network", "node", node_name, "application", "DatabaseClient", "configure", config] - - -class ConfigureRansomwareScriptAction(AbstractAction): - """Action which sets config parameters for a ransomware script on a node.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this option.""" - - model_config = ConfigDict(extra="forbid") - server_ip_address: Optional[str] = None - server_password: Optional[str] = None - payload: Optional[str] = None - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, config: Dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema - return ["network", "node", node_name, "application", "RansomwareScript", "configure", config] - - -class ConfigureDoSBotAction(AbstractAction): - """Action which sets config parameters for a DoS bot on a node.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this action.""" - - model_config = ConfigDict(extra="forbid") - target_ip_address: Optional[str] = None - target_port: Optional[str] = None - payload: Optional[str] = None - repeat: Optional[bool] = None - port_scan_p_of_success: Optional[float] = None - dos_intensity: Optional[float] = None - max_sessions: Optional[int] = None - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, config: Dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - self._Opts.model_validate(config) # check that options adhere to schema - return ["network", "node", node_name, "application", "DoSBot", "configure", config] - - -class NodeApplicationRemoveAction(AbstractAction): - """Action which removes/uninstalls an application.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes} - - def form_request(self, node_id: int, application_name: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "software_manager", "application", "uninstall", application_name] - - -class NodeFolderAbstractAction(AbstractAction): - """ - Base class for folder actions. - - Any action which applies to a folder and uses node_id and folder_id as its only two parameters can inherit from - this base class. - """ - - @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int, folder_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id) - if node_name is None or folder_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "folder", folder_name, self.verb] - - -class NodeFolderScanAction(NodeFolderAbstractAction): - """Action which scans a folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "scan" - - -class NodeFolderCheckhashAction(NodeFolderAbstractAction): - """Action which checks the hash of a folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "checkhash" - - -class NodeFolderRepairAction(NodeFolderAbstractAction): - """Action which repairs a folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "repair" - - -class NodeFolderRestoreAction(NodeFolderAbstractAction): - """Action which restores a folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "restore" - - -class NodeFileCreateAction(AbstractAction): - """Action which creates a new file in a given folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "create" - - def form_request( - self, node_id: int, folder_name: str, file_name: str, force: Optional[bool] = False - ) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None or folder_name is None or file_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name, force] - - -class NodeFolderCreateAction(AbstractAction): - """Action which creates a new folder.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "create" - - def form_request(self, node_id: int, folder_name: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None or folder_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "create", "folder", folder_name] - - -class NodeFileAbstractAction(AbstractAction): - """Abstract base class for file actions. - - Any action which applies to a file and uses node_id, folder_id, and file_id as its only three parameters can inherit - from this base class. - """ - - @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders, "file_id": num_files} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int, folder_id: int, file_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id) - file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id) - if node_name is None or folder_name is None or file_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "folder", folder_name, "file", file_name, self.verb] - - -class NodeFileScanAction(NodeFileAbstractAction): - """Action which scans a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "scan" - - -class NodeFileCheckhashAction(NodeFileAbstractAction): - """Action which checks the hash of a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "checkhash" - - -class NodeFileDeleteAction(NodeFileAbstractAction): - """Action which deletes a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "delete" - - def form_request(self, node_id: int, folder_id: int, file_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id) - file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id) - if node_name is None or folder_name is None or file_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "delete", "file", folder_name, file_name] - - -class NodeFileRepairAction(NodeFileAbstractAction): - """Action which repairs a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "repair" - - -class NodeFileRestoreAction(NodeFileAbstractAction): - """Action which restores a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "restore" - - -class NodeFileCorruptAction(NodeFileAbstractAction): - """Action which corrupts a file.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) - self.verb: str = "corrupt" - - -class NodeFileAccessAction(AbstractAction): - """Action which increases a file's access count.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) - self.verb: str = "access" - - def form_request(self, node_id: int, folder_name: str, file_name: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None or folder_name is None or file_name is None: - return ["do_nothing"] - return ["network", "node", node_name, "file_system", "access", folder_name, file_name] - - -class NodeAbstractAction(AbstractAction): - """ - Abstract base class for node actions. - - Any action which applies to a node and uses node_id as its only parameter can inherit from this base class. - """ - - @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - return ["network", "node", node_name, self.verb] - - -class NodeOSScanAction(NodeAbstractAction): - """Action which scans a node's OS.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes) - self.verb: str = "scan" - - -class NodeShutdownAction(NodeAbstractAction): - """Action which shuts down a node.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes) - self.verb: str = "shutdown" - - -class NodeStartupAction(NodeAbstractAction): - """Action which starts up a node.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes) - self.verb: str = "startup" - - -class NodeResetAction(NodeAbstractAction): - """Action which resets a node.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes) - self.verb: str = "reset" - - -class RouterACLAddRuleAction(AbstractAction): - """Action which adds a rule to a router's ACL.""" - - class ACLRuleOptions(BaseModel): - """Validator for ACL_ADD_RULE options.""" - - target_router: str - """On which router to add the rule, must be specified.""" - position: int - """At what position to add the rule, must be specified.""" - permission: Literal[1, 2] - """Whether to allow or deny traffic, must be specified. 1 = PERMIT, 2 = DENY.""" - source_ip_id: int = Field(default=1, ge=1) - """Rule source IP address. By default, all ip addresses.""" - source_wildcard_id: int = Field(default=0, ge=0) - """Rule source IP wildcard. By default, use the wildcard at index 0 from action manager.""" - source_port_id: int = Field(default=1, ge=1) - """Rule source port. By default, all source ports.""" - dest_ip_id: int = Field(default=1, ge=1) - """Rule destination IP address. By default, all ip addresses.""" - dest_wildcard_id: int = Field(default=0, ge=0) - """Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager.""" - dest_port_id: int = Field(default=1, ge=1) - """Rule destination port. By default, all destination ports.""" - protocol_id: int = Field(default=1, ge=1) - """Rule protocol. By default, all protocols.""" - - @field_validator( - "source_ip_id", - "source_port_id", - "source_wildcard_id", - "dest_ip_id", - "dest_port_id", - "dest_wildcard_id", - "protocol_id", - mode="before", - ) - @classmethod - def not_none(cls, v: str, info: ValidationInfo) -> int: - """If None is passed, use the default value instead.""" - if v is None: - return cls.model_fields[info.field_name].default - return v - - def __init__( - self, - manager: "ActionManager", - max_acl_rules: int, - num_ips: int, - num_ports: int, - num_protocols: int, - **kwargs, - ) -> None: - """Init method for RouterACLAddRuleAction. - - :param manager: Reference to the ActionManager which created this action. - :type manager: ActionManager - :param max_acl_rules: Maximum number of ACL rules that can be added to the router. - :type max_acl_rules: int - :param num_ips: Number of IP addresses in the simulation. - :type num_ips: int - :param num_ports: Number of ports in the simulation. - :type num_ports: int - :param num_protocols: Number of protocols in the simulation. - :type num_protocols: int - """ - super().__init__(manager=manager) - num_permissions = 3 - self.shape: Dict[str, int] = { - "position": max_acl_rules, - "permission": num_permissions, - "source_ip_id": num_ips, - "dest_ip_id": num_ips, - "source_port_id": num_ports, - "dest_port_id": num_ports, - "protocol_id": num_protocols, - } - - def form_request( - self, - target_router: str, - position: int, - permission: int, - source_ip_id: int, - source_wildcard_id: int, - dest_ip_id: int, - dest_wildcard_id: int, - source_port_id: int, - dest_port_id: int, - protocol_id: int, - ) -> List[str]: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - # Validate incoming data. - parsed_options = RouterACLAddRuleAction.ACLRuleOptions( - target_router=target_router, - position=position, - permission=permission, - source_ip_id=source_ip_id, - source_wildcard_id=source_wildcard_id, - dest_ip_id=dest_ip_id, - dest_wildcard_id=dest_wildcard_id, - source_port_id=source_port_id, - dest_port_id=dest_port_id, - protocol_id=protocol_id, - ) - if parsed_options.permission == 1: - permission_str = "PERMIT" - elif parsed_options.permission == 2: - permission_str = "DENY" - else: - _LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.") - - if parsed_options.protocol_id == 1: - protocol = "ALL" - else: - protocol = self.manager.get_internet_protocol_by_idx(parsed_options.protocol_id - 2) - # subtract 2 to account for UNUSED=0 and ALL=1. - - if parsed_options.source_ip_id == 1: - src_ip = "ALL" - else: - src_ip = self.manager.get_ip_address_by_idx(parsed_options.source_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - src_wildcard = self.manager.get_wildcard_by_idx(parsed_options.source_wildcard_id) - - if parsed_options.source_port_id == 1: - src_port = "ALL" - else: - src_port = self.manager.get_port_by_idx(parsed_options.source_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - if parsed_options.dest_ip_id == 1: - dst_ip = "ALL" - else: - dst_ip = self.manager.get_ip_address_by_idx(parsed_options.dest_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - dst_wildcard = self.manager.get_wildcard_by_idx(parsed_options.dest_wildcard_id) - - if parsed_options.dest_port_id == 1: - dst_port = "ALL" - else: - dst_port = self.manager.get_port_by_idx(parsed_options.dest_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - return [ - "network", - "node", - target_router, - "acl", - "add_rule", - permission_str, - protocol, - str(src_ip), - src_wildcard, - src_port, - str(dst_ip), - dst_wildcard, - dst_port, - position, - ] - - -class RouterACLRemoveRuleAction(AbstractAction): - """Action which removes a rule from a router's ACL.""" - - def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None: - """Init method for RouterACLRemoveRuleAction. - - :param manager: Reference to the ActionManager which created this action. - :type manager: ActionManager - :param max_acl_rules: Maximum number of ACL rules that can be added to the router. - :type max_acl_rules: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"position": max_acl_rules} - - def form_request(self, target_router: str, position: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return ["network", "node", target_router, "acl", "remove_rule", position] - - -class FirewallACLAddRuleAction(AbstractAction): - """Action which adds a rule to a firewall port's ACL.""" - - def __init__( - self, - manager: "ActionManager", - max_acl_rules: int, - num_ips: int, - num_ports: int, - num_protocols: int, - **kwargs, - ) -> None: - """Init method for FirewallACLAddRuleAction. - - :param manager: Reference to the ActionManager which created this action. - :type manager: ActionManager - :param max_acl_rules: Maximum number of ACL rules that can be added to the router. - :type max_acl_rules: int - :param num_ips: Number of IP addresses in the simulation. - :type num_ips: int - :param num_ports: Number of ports in the simulation. - :type num_ports: int - :param num_protocols: Number of protocols in the simulation. - :type num_protocols: int - """ - super().__init__(manager=manager) - num_permissions = 3 - self.shape: Dict[str, int] = { - "position": max_acl_rules, - "permission": num_permissions, - "source_ip_id": num_ips, - "dest_ip_id": num_ips, - "source_port_id": num_ports, - "dest_port_id": num_ports, - "protocol_id": num_protocols, - } - - def form_request( - self, - target_firewall_nodename: str, - firewall_port_name: str, - firewall_port_direction: str, - position: int, - permission: int, - source_ip_id: int, - source_wildcard_id: int, - dest_ip_id: int, - dest_wildcard_id: int, - source_port_id: int, - dest_port_id: int, - protocol_id: int, - ) -> List[str]: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - if permission == 0: - permission_str = "UNUSED" - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS - elif permission == 1: - permission_str = "PERMIT" - elif permission == 2: - permission_str = "DENY" - else: - _LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.") - - if protocol_id == 0: - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS - - if protocol_id == 1: - protocol = "ALL" - else: - protocol = self.manager.get_internet_protocol_by_idx(protocol_id - 2) - # subtract 2 to account for UNUSED=0 and ALL=1. - - if source_ip_id == 0: - return ["do_nothing"] # invalid formulation - elif source_ip_id == 1: - src_ip = "ALL" - else: - src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - if source_port_id == 0: - return ["do_nothing"] # invalid formulation - elif source_port_id == 1: - src_port = "ALL" - else: - src_port = self.manager.get_port_by_idx(source_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - if dest_ip_id == 0: - return ["do_nothing"] # invalid formulation - elif dest_ip_id == 1: - dst_ip = "ALL" - else: - dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - if dest_port_id == 0: - return ["do_nothing"] # invalid formulation - elif dest_port_id == 1: - dst_port = "ALL" - else: - dst_port = self.manager.get_port_by_idx(dest_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id) - dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id) - - return [ - "network", - "node", - target_firewall_nodename, - firewall_port_name, - firewall_port_direction, - "acl", - "add_rule", - permission_str, - protocol, - str(src_ip), - src_wildcard, - src_port, - str(dst_ip), - dst_wildcard, - dst_port, - position, - ] - - -class FirewallACLRemoveRuleAction(AbstractAction): - """Action which removes a rule from a firewall port's ACL.""" - - def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None: - """Init method for RouterACLRemoveRuleAction. - - :param manager: Reference to the ActionManager which created this action. - :type manager: ActionManager - :param max_acl_rules: Maximum number of ACL rules that can be added to the router. - :type max_acl_rules: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"position": max_acl_rules} - - def form_request( - self, target_firewall_nodename: str, firewall_port_name: str, firewall_port_direction: str, position: int - ) -> List[str]: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [ - "network", - "node", - target_firewall_nodename, - firewall_port_name, - firewall_port_direction, - "acl", - "remove_rule", - position, - ] - - -class HostNICAbstractAction(AbstractAction): - """ - Abstract base class for NIC actions. - - Any action which applies to a NIC and uses node_id and nic_id as its only two parameters can inherit from this base - class. - """ - - def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: - """Init method for HostNICAbstractAction. - - :param manager: Reference to the ActionManager which created this action. - :type manager: ActionManager - :param num_nodes: Number of nodes in the simulation. - :type num_nodes: int - :param max_nics_per_node: Maximum number of NICs per node. - :type max_nics_per_node: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} - self.verb: str # define but don't initialise: defends against children classes not defining this - - def form_request(self, node_id: int, nic_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_idx=node_id) - nic_num = self.manager.get_nic_num_by_idx(node_idx=node_id, nic_idx=nic_id) - if node_name is None or nic_num is None: - return ["do_nothing"] - return ["network", "node", node_name, "network_interface", nic_num, self.verb] - - -class HostNICEnableAction(HostNICAbstractAction): - """Action which enables a NIC.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) - self.verb: str = "enable" - - -class HostNICDisableAction(HostNICAbstractAction): - """Action which disables a NIC.""" - - def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: - super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) - self.verb: str = "disable" - - -class NetworkPortEnableAction(AbstractAction): - """Action which enables are port on a router or a firewall.""" - - def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None: - """Init method for NetworkPortEnableAction. - - :param max_nics_per_node: Maximum number of NICs per node. - :type max_nics_per_node: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"port_id": max_nics_per_node} - - def form_request(self, target_nodename: str, port_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - if target_nodename is None or port_id is None: - return ["do_nothing"] - return ["network", "node", target_nodename, "network_interface", port_id, "enable"] - - -class NetworkPortDisableAction(AbstractAction): - """Action which disables are port on a router or a firewall.""" - - def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None: - """Init method for NetworkPortDisableAction. - - :param max_nics_per_node: Maximum number of NICs per node. - :type max_nics_per_node: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"port_id": max_nics_per_node} - - def form_request(self, target_nodename: str, port_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - if target_nodename is None or port_id is None: - return ["do_nothing"] - return ["network", "node", target_nodename, "network_interface", port_id, "disable"] - - -class NodeNMAPPingScanAction(AbstractAction): - """Action which performs an NMAP ping scan.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request( - self, source_node: str, target_ip_address: Union[str, List[str]], show: Optional[bool] = False - ) -> List[str]: # noqa - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [ - "network", - "node", - source_node, - "application", - "NMAP", - "ping_scan", - {"target_ip_address": target_ip_address, "show": show}, - ] - - -class NodeNMAPPortScanAction(AbstractAction): - """Action which performs an NMAP port scan.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request( - self, - source_node: str, - target_ip_address: Union[str, List[str]], - target_protocol: Optional[Union[str, List[str]]] = None, - target_port: Optional[Union[str, List[str]]] = None, - show: Optional[bool] = False, - ) -> List[str]: # noqa - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [ - "network", - "node", - source_node, - "application", - "NMAP", - "port_scan", - { - "target_ip_address": target_ip_address, - "target_port": target_port, - "target_protocol": target_protocol, - "show": show, - }, - ] - - -class NodeNetworkServiceReconAction(AbstractAction): - """Action which performs an NMAP network service recon (ping scan followed by port scan).""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request( - self, - source_node: str, - target_ip_address: Union[str, List[str]], - target_protocol: Optional[Union[str, List[str]]] = None, - target_port: Optional[Union[str, List[str]]] = None, - show: Optional[bool] = False, - ) -> List[str]: # noqa - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [ - "network", - "node", - source_node, - "application", - "NMAP", - "network_service_recon", - { - "target_ip_address": target_ip_address, - "target_port": target_port, - "target_protocol": target_protocol, - "show": show, - }, - ] - - -class ConfigureC2BeaconAction(AbstractAction): - """Action which configures a C2 Beacon based on the parameters given.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this action.""" - - c2_server_ip_address: str - keep_alive_frequency: int = Field(default=5, ge=1) - masquerade_protocol: str = Field(default="TCP") - masquerade_port: str = Field(default="HTTP") - - @field_validator( - "c2_server_ip_address", - "keep_alive_frequency", - "masquerade_protocol", - "masquerade_port", - mode="before", - ) - @classmethod - def not_none(cls, v: str, info: ValidationInfo) -> int: - """If None is passed, use the default value instead.""" - if v is None: - return cls.model_fields[info.field_name].default - return v - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, config: Dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - config = ConfigureC2BeaconAction._Opts( - c2_server_ip_address=config["c2_server_ip_address"], - keep_alive_frequency=config["keep_alive_frequency"], - masquerade_port=config["masquerade_port"], - masquerade_protocol=config["masquerade_protocol"], - ) - - ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema - - return ["network", "node", node_name, "application", "C2Beacon", "configure", config.__dict__] - - -class NodeAccountsChangePasswordAction(AbstractAction): - """Action which changes the password for a user.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: str, username: str, current_password: str, new_password: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - return [ - "network", - "node", - node_name, - "service", - "UserManager", - "change_password", - username, - current_password, - new_password, - ] - - -class NodeSessionsRemoteLoginAction(AbstractAction): - """Action which performs a remote session login.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - return [ - "network", - "node", - node_name, - "service", - "Terminal", - "ssh_to_remote", - username, - password, - remote_ip, - ] - - -class NodeSessionsRemoteLogoutAction(AbstractAction): - """Action which performs a remote session logout.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: str, remote_ip: str) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip] - - -class RansomwareConfigureC2ServerAction(AbstractAction): - """Action which sends a command from the C2 Server to the C2 Beacon which configures a local RansomwareScript.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, config: Dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - # Using the ransomware scripts model to validate. - ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema - return ["network", "node", node_name, "application", "C2Server", "ransomware_configure", config] - - -class RansomwareLaunchC2ServerAction(AbstractAction): - """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - # This action currently doesn't require any further configuration options. - return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"] - - -class ExfiltrationC2ServerAction(AbstractAction): - """Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this action.""" - - username: Optional[str] - password: Optional[str] - target_ip_address: str - target_file_name: str - target_folder_name: str - exfiltration_folder_name: Optional[str] - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request( - self, - node_id: int, - account: dict, - target_ip_address: str, - target_file_name: str, - target_folder_name: str, - exfiltration_folder_name: Optional[str], - ) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - - command_model = { - "target_file_name": target_file_name, - "target_folder_name": target_folder_name, - "exfiltration_folder_name": exfiltration_folder_name, - "target_ip_address": target_ip_address, - "username": account["username"], - "password": account["password"], - } - ExfiltrationC2ServerAction._Opts.model_validate(command_model) - return ["network", "node", node_name, "application", "C2Server", "exfiltrate", command_model] - - -class NodeSendRemoteCommandAction(AbstractAction): - """Action which sends a terminal command to a remote node via SSH.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, remote_ip: str, command: RequestFormat) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - return [ - "network", - "node", - node_name, - "service", - "Terminal", - "send_remote_command", - remote_ip, - {"command": command}, - ] - - -class TerminalC2ServerAction(AbstractAction): - """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" - - class _Opts(BaseModel): - """Schema for options that can be passed to this action.""" - - commands: Union[List[RequestFormat], RequestFormat] - ip_address: Optional[str] - username: Optional[str] - password: Optional[str] - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int, commands: List, ip_address: Optional[str], account: dict) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - - command_model = { - "commands": commands, - "ip_address": ip_address, - "username": account["username"], - "password": account["password"], - } - - TerminalC2ServerAction._Opts.model_validate(command_model) - return ["network", "node", node_name, "application", "C2Server", "terminal_command", command_model] - - -class RansomwareLaunchC2ServerAction(AbstractAction): - """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" - - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, node_id: int) -> RequestFormat: - """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: - return ["do_nothing"] - # This action currently doesn't require any further configuration options. - return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"] - - -class ActionManager: - """Class which manages the action space for an agent.""" - - act_class_identifiers: Dict[str, type] = { - "DONOTHING": DoNothingAction, - "NODE_SERVICE_SCAN": NodeServiceScanAction, - "NODE_SERVICE_STOP": NodeServiceStopAction, - "NODE_SERVICE_START": NodeServiceStartAction, - "NODE_SERVICE_PAUSE": NodeServicePauseAction, - "NODE_SERVICE_RESUME": NodeServiceResumeAction, - "NODE_SERVICE_RESTART": NodeServiceRestartAction, - "NODE_SERVICE_DISABLE": NodeServiceDisableAction, - "NODE_SERVICE_ENABLE": NodeServiceEnableAction, - "NODE_SERVICE_FIX": NodeServiceFixAction, - "NODE_APPLICATION_EXECUTE": NodeApplicationExecuteAction, - "NODE_APPLICATION_SCAN": NodeApplicationScanAction, - "NODE_APPLICATION_CLOSE": NodeApplicationCloseAction, - "NODE_APPLICATION_FIX": NodeApplicationFixAction, - "NODE_APPLICATION_INSTALL": NodeApplicationInstallAction, - "NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction, - "NODE_FILE_SCAN": NodeFileScanAction, - "NODE_FILE_CREATE": NodeFileCreateAction, - "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, - "NODE_FILE_DELETE": NodeFileDeleteAction, - "NODE_FILE_REPAIR": NodeFileRepairAction, - "NODE_FILE_RESTORE": NodeFileRestoreAction, - "NODE_FILE_CORRUPT": NodeFileCorruptAction, - "NODE_FILE_ACCESS": NodeFileAccessAction, - "NODE_FOLDER_CREATE": NodeFolderCreateAction, - "NODE_FOLDER_SCAN": NodeFolderScanAction, - "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, - "NODE_FOLDER_REPAIR": NodeFolderRepairAction, - "NODE_FOLDER_RESTORE": NodeFolderRestoreAction, - "NODE_OS_SCAN": NodeOSScanAction, - "NODE_SHUTDOWN": NodeShutdownAction, - "NODE_STARTUP": NodeStartupAction, - "NODE_RESET": NodeResetAction, - "ROUTER_ACL_ADDRULE": RouterACLAddRuleAction, - "ROUTER_ACL_REMOVERULE": RouterACLRemoveRuleAction, - "FIREWALL_ACL_ADDRULE": FirewallACLAddRuleAction, - "FIREWALL_ACL_REMOVERULE": FirewallACLRemoveRuleAction, - "HOST_NIC_ENABLE": HostNICEnableAction, - "HOST_NIC_DISABLE": HostNICDisableAction, - "NETWORK_PORT_ENABLE": NetworkPortEnableAction, - "NETWORK_PORT_DISABLE": NetworkPortDisableAction, - "NODE_NMAP_PING_SCAN": NodeNMAPPingScanAction, - "NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction, - "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction, - "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, - "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, - "CONFIGURE_DOSBOT": ConfigureDoSBotAction, - "CONFIGURE_C2_BEACON": ConfigureC2BeaconAction, - "C2_SERVER_RANSOMWARE_LAUNCH": RansomwareLaunchC2ServerAction, - "C2_SERVER_RANSOMWARE_CONFIGURE": RansomwareConfigureC2ServerAction, - "C2_SERVER_TERMINAL_COMMAND": TerminalC2ServerAction, - "C2_SERVER_DATA_EXFILTRATE": ExfiltrationC2ServerAction, - "NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction, - "SSH_TO_REMOTE": NodeSessionsRemoteLoginAction, - "SESSIONS_REMOTE_LOGOFF": NodeSessionsRemoteLogoutAction, - "NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction, - } - """Dictionary which maps action type strings to the corresponding action class.""" - - def __init__( - self, - 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 - max_files_per_folder: int = 2, # allows calculating shape - max_services_per_node: int = 2, # allows calculating shape - max_applications_per_node: int = 2, # allows calculating shape - max_nics_per_node: int = 8, # allows calculating shape - 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_list: List[str] = [], # to allow us to map an index to an ip address. - wildcard_list: List[str] = [], # to allow mapping from wildcard index to - act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions - ) -> None: - """Init method for ActionManager. - - :param game: Reference to the game to which the agent belongs. - :type game: PrimaiteGame - :param actions: List of action specs which should be made available to the agent. The keys of each spec are: - 'type' and 'options' for passing any options to the action class's init method - :type actions: List[dict] - :param nodes: Extra configuration for each node. - :type nodes: List[Dict] - :param max_folders_per_node: Maximum number of folders per node. Used for calculating action shape. - :type max_folders_per_node: int - :param max_files_per_folder: Maximum number of files per folder. Used for calculating action shape. - :type max_files_per_folder: int - :param max_services_per_node: Maximum number of services per node. Used for calculating action shape. - :type max_services_per_node: int - :param max_nics_per_node: Maximum number of NICs per node. Used for calculating action shape. - :type max_nics_per_node: int - :param max_acl_rules: Maximum number of ACL rules per router. Used for calculating action shape. - :type max_acl_rules: int - :param protocols: List of protocols that are available in the simulation. Used for calculating action shape. - :type protocols: List[str] - :param ports: List of ports that are available in the simulation. Used for calculating action shape. - :type ports: List[str] - :param ip_list: List of IP addresses that known to this agent. Used for calculating action shape. - :type ip_list: Optional[List[str]] - :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.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]] = [] - """ - List of applications per node. The list order gives the two-index mapping between (node_id, app_id) to app name. - The first index corresponds to node id, the second index is the app id on that particular node. - For instance, self.application_names[0][2] is the name of the third application on the first node. - """ - self.service_names: List[List[str]] = [] - """ - List of services per node. The list order gives the two-index mapping between (node_id, svc_id) to svc name. - The first index corresponds to node id, the second index is the service id on that particular node. - For instance, self.service_names[0][2] is the name of the third service on the first node. - """ - self.folder_names: List[List[str]] = [] - """ - List of folders per node. The list order gives the two-index mapping between (node_id, folder_id) to folder - name. The first index corresponds to node id, the second index is the folder id on that particular node. - For instance, self.folder_names[0][2] is the name of the third folder on the first node. - """ - self.file_names: List[List[List[str]]] = [] - """ - List of files per folder per node. The list order gives the three-index mapping between - (node_id, folder_id, file_id) to file name. The first index corresponds to node id, the second index is the - folder id on that particular node, and the third index is the file id in that particular folder. - For instance, self.file_names[0][2][1] is the name of the second file in the third folder on the first node. - """ - - # Populate lists of apps, services, files, folders, etc on nodes. - for node in nodes: - app_list = [a["application_name"] for a in node.get("applications", [])] - while len(app_list) < max_applications_per_node: - app_list.append(None) - self.application_names.append(app_list) - - svc_list = [s["service_name"] for s in node.get("services", [])] - while len(svc_list) < max_services_per_node: - svc_list.append(None) - self.service_names.append(svc_list) - - folder_list = [f["folder_name"] for f in node.get("folders", [])] - while len(folder_list) < max_folders_per_node: - folder_list.append(None) - self.folder_names.append(folder_list) - - file_sublist = [] - for folder in node.get("folders", [{"files": []}]): - file_list = [f["file_name"] for f in folder.get("files", [])] - while len(file_list) < max_files_per_folder: - file_list.append(None) - file_sublist.append(file_list) - while len(file_sublist) < max_folders_per_node: - file_sublist.append([None] * max_files_per_folder) - self.file_names.append(file_sublist) - self.protocols: List[str] = protocols - self.ports: List[str] = ports - - self.ip_address_list: List[str] = ip_list - self.wildcard_list: List[str] = wildcard_list - if self.wildcard_list == []: - self.wildcard_list = ["NONE"] - # action_args are settings which are applied to the action space as a whole. - global_action_args = { - "num_nodes": len(self.node_names), - "num_folders": max_folders_per_node, - "num_files": max_files_per_folder, - "num_services": max_services_per_node, - "num_applications": max_applications_per_node, - "num_nics": max_nics_per_node, - "num_acl_rules": max_acl_rules, - "num_protocols": len(self.protocols), - "num_ports": len(self.protocols), - "num_ips": len(self.ip_address_list), - "max_acl_rules": max_acl_rules, - "max_nics_per_node": max_nics_per_node, - } - self.actions: Dict[str, AbstractAction] = {} - for act_spec in actions: - # each action is provided into the action space config like this: - # - type: ACTION_TYPE - # options: - # option_1: value1 - # option_2: value2 - # where `type` decides which AbstractAction subclass should be used - # 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.action_map: Dict[int, Tuple[str, Dict]] = {} - """ - Action mapping that converts an integer to a specific action and parameter choice. - - For example : - {0: ("NODE_SERVICE_SCAN", {node_id:0, service_id:2})} - """ - if act_map is None: - # raise RuntimeError("Action map must be specified in the config file.") - pass - else: - self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()} - # make sure all numbers between 0 and N are represented as dict keys in action map - assert all([i in self.action_map.keys() for i in range(len(self.action_map))]) - - def _enumerate_actions( - self, - ) -> Dict[int, Tuple[str, Dict]]: - """Generate a list of all the possible actions that could be taken. - - This enumerates all actions all combinations of parameters you could choose for those actions. The output - of this function is intended to populate the self.action_map parameter in the situation where the user provides - a list of action types, but doesn't specify any subset of actions that should be made available to the agent. - - The enumeration relies on the Actions' `shape` attribute. - - :return: An action map maps consecutive integers to a combination of Action type and parameter choices. - An example output could be: - {0: ("DONOTHING", {'dummy': 0}), - 1: ("NODE_OS_SCAN", {'node_id': 0}), - 2: ("NODE_OS_SCAN", {'node_id': 1}), - 3: ("NODE_FOLDER_SCAN", {'node_id:0, folder_id:0}), - ... #etc... - } - :rtype: Dict[int, Tuple[AbstractAction, Dict]] - """ - all_action_possibilities = [] - for act_name, action in self.actions.items(): - param_names = list(action.shape.keys()) - num_possibilities = list(action.shape.values()) - possibilities = [range(n) for n in num_possibilities] - - param_combinations = list(itertools.product(*possibilities)) - all_action_possibilities.extend( - [ - (act_name, {param_names[i]: param_combinations[j][i] for i in range(len(param_names))}) - for j in range(len(param_combinations)) - ] - ) - - return {i: p for i, p in enumerate(all_action_possibilities)} - - def get_action(self, action: int) -> Tuple[str, Dict]: - """Produce action in CAOS format.""" - """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" - """The CAOS format is basically a action identifier, followed by parameters stored in a dictionary""" - act_identifier, act_options = self.action_map[action] - return act_identifier, act_options - - def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat: - """Take action in CAOS format and use the execution definition to change it into PrimAITE request format.""" - act_obj = self.actions[action_identifier] - return act_obj.form_request(**action_options) - - @property - def space(self) -> spaces.Space: - """Return the gymnasium action space for this agent.""" - return spaces.Discrete(len(self.action_map)) - - def get_node_name_by_idx(self, node_idx: int) -> str: - """ - Get the node name corresponding to the given index. - - :param node_idx: The index of the node to retrieve. - :type node_idx: int - :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]: - """ - Get the folder name corresponding to the given node and folder indices. - - :param node_idx: The index of the node. - :type node_idx: int - :param folder_idx: The index of the folder on the node. - :type folder_idx: int - :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]: - """Get the file name corresponding to the given node, folder, and file indices. - - :param node_idx: The index of the node. - :type node_idx: int - :param folder_idx: The index of the folder on the node. - :type folder_idx: int - :param file_idx: The index of the file in the folder. - :type file_idx: int - :return: The name of the file. Or None if the node has fewer folders than the given index, or the folder has - 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]: - """Get the service name corresponding to the given node and service indices. - - :param node_idx: The index of the node. - :type node_idx: int - :param service_idx: The index of the service on the node. - :type service_idx: int - :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]: - """Get the application name corresponding to the given node and service indices. - - :param node_idx: The index of the node. - :type node_idx: int - :param application_idx: The index of the service on the node. - :type application_idx: int - :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: - """Get the internet protocol corresponding to the given index. - - :param protocol_idx: The index of the protocol to retrieve. - :type protocol_idx: int - :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: - """ - Get the IP address corresponding to the given index. - - :param ip_idx: The index of the IP address to retrieve. - :type ip_idx: int - :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_wildcard_by_idx(self, wildcard_idx: int) -> str: - """ - Get the IP wildcard corresponding to the given index. - - :param ip_idx: The index of the IP wildcard to retrieve. - :type ip_idx: int - :return: The wildcard address. - :rtype: str - """ - if wildcard_idx >= len(self.wildcard_list): - msg = ( - f"Error: agent attempted to perform an action on ip wildcard {wildcard_idx} but this" - f" is out of range for its action space. Wildcard list: {self.wildcard_list}" - ) - _LOGGER.error(msg) - raise RuntimeError(msg) - return self.wildcard_list[wildcard_idx] - - def get_port_by_idx(self, port_idx: int) -> str: - """ - Get the port corresponding to the given index. - - :param port_idx: The index of the port to retrieve. - :type port_idx: int - :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: - """ - Get the NIC number corresponding to the given node and NIC indices. - - :param node_idx: The index of the node. - :type node_idx: int - :param nic_idx: The index of the NIC on the node. - :type nic_idx: int - :return: The NIC number. - :rtype: int - """ - return nic_idx + 1 - - @classmethod - def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager": - """ - Construct an ActionManager from a config definition. - - The action space config supports the following three sections: - 1. ``action_list`` - ``action_list`` contains a list action components which need to be included in the action space. - Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options - which will be passed to the action class's __init__ method during initialisation. - 2. ``action_map`` - Since the agent uses a discrete action space which acts as a flattened version of the component-based - action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful - action and values of parameters. For example action 0 can correspond to do nothing, action 1 can - correspond to "NODE_SERVICE_SCAN" with ``node_id=1`` and ``service_id=1``, action 2 can be " - 3. ``options`` - ``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method. - These options are used to calculate the shape of the action space, and to provide additional information - to the ActionManager which is required to convert the agent's action choice into a CAOS request. - - :param game: The Primaite Game to which the agent belongs. - :type game: PrimaiteGame - :param cfg: The action space config. - :type cfg: Dict - :return: The constructed ActionManager. - :rtype: ActionManager - """ - if "ip_list" not in cfg["options"]: - cfg["options"]["ip_list"] = [] - - obj = cls( - actions=cfg["action_list"], - **cfg["options"], - protocols=game.options.protocols, - ports=game.options.ports, - act_map=cfg.get("action_map"), - ) - - return obj diff --git a/src/primaite/game/agent/actions/__init__.py b/src/primaite/game/agent/actions/__init__.py new file mode 100644 index 00000000..8517ded8 --- /dev/null +++ b/src/primaite/game/agent/actions/__init__.py @@ -0,0 +1,33 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + +from primaite.game.agent.actions import ( + abstract, + acl, + application, + file, + folder, + host_nic, + manager, + network, + node, + service, + session, + software, +) +from primaite.game.agent.actions.manager import ActionManager + +__all__ = ( + "abstract", + "acl", + "application", + "software", + "file", + "folder", + "host_nic", + "manager", + "network", + "node", + "service", + "session", + "ActionManager", +) diff --git a/src/primaite/game/agent/actions/abstract.py b/src/primaite/game/agent/actions/abstract.py new file mode 100644 index 00000000..1cda4360 --- /dev/null +++ b/src/primaite/game/agent/actions/abstract.py @@ -0,0 +1,36 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from __future__ import annotations + +from abc import ABC +from typing import Any, ClassVar, Dict, Optional, Type + +from pydantic import BaseModel, ConfigDict + +from primaite.interface.request import RequestFormat + + +class AbstractAction(BaseModel, ABC): + """Base class for actions.""" + + config: "AbstractAction.ConfigSchema" + + class ConfigSchema(BaseModel, ABC): + """Base configuration schema for Actions.""" + + model_config = ConfigDict(extra="forbid") + type: str = "" + + _registry: ClassVar[Dict[str, Type[AbstractAction]]] = {} + + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if identifier is None: + return + if identifier in cls._registry: + raise ValueError(f"Cannot create new action under reserved name {identifier}") + cls._registry[identifier] = cls + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + pass diff --git a/src/primaite/game/agent/actions/acl.py b/src/primaite/game/agent/actions/acl.py new file mode 100644 index 00000000..3341868f --- /dev/null +++ b/src/primaite/game/agent/actions/acl.py @@ -0,0 +1,157 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from __future__ import annotations + +from abc import ABC +from typing import Literal, Union + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat +from primaite.utils.validation.ip_protocol import IPProtocol +from primaite.utils.validation.ipv4_address import IPV4Address +from primaite.utils.validation.port import Port + +__all__ = ( + "RouterACLAddRuleAction", + "RouterACLRemoveRuleAction", + "FirewallACLAddRuleAction", + "FirewallACLRemoveRuleAction", +) + + +class ACLAddRuleAbstractAction(AbstractAction, ABC): + """Base abstract class for ACL add rule actions.""" + + config: ConfigSchema = "ACLAddRuleAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration Schema base for ACL add rule abstract actions.""" + + src_ip: IPV4Address + protocol_name: Union[IPProtocol, Literal["ALL"]] + permission: Literal["PERMIT", "DENY"] + position: int + dst_ip: Union[IPV4Address, Literal["ALL"]] + src_port: Union[Port, Literal["ALL"]] + dst_port: Union[Port, Literal["ALL"]] + src_wildcard: Union[IPV4Address, Literal["NONE"]] + dst_wildcard: Union[IPV4Address, Literal["NONE"]] + + +class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_abstract_action"): + """Base abstract class for ACL remove rule actions.""" + + config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration Schema base for ACL remove rule abstract actions.""" + + position: int + + +class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_add_rule"): + """Action which adds a rule to a router's ACL.""" + + config: "RouterACLAddRuleAction.ConfigSchema" + + class ConfigSchema(ACLAddRuleAbstractAction.ConfigSchema): + """Configuration Schema for RouterACLAddRuleAction.""" + + target_router: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.target_router, + "acl", + "add_rule", + config.permission, + config.protocol_name, + str(config.src_ip), + str(config.src_wildcard), + config.src_port, + str(config.dst_ip), + str(config.dst_wildcard), + config.dst_port, + config.position, + ] + + +class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_acl_remove_rule"): + """Action which removes a rule from a router's ACL.""" + + config: "RouterACLRemoveRuleAction.ConfigSchema" + + class ConfigSchema(ACLRemoveRuleAbstractAction.ConfigSchema): + """Configuration schema for RouterACLRemoveRuleAction.""" + + target_router: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return ["network", "node", config.target_router, "acl", "remove_rule", config.position] + + +class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_acl_add_rule"): + """Action which adds a rule to a firewall port's ACL.""" + + config: "FirewallACLAddRuleAction.ConfigSchema" + + class ConfigSchema(ACLAddRuleAbstractAction.ConfigSchema): + """Configuration schema for FirewallACLAddRuleAction.""" + + target_firewall_nodename: str + firewall_port_name: str + firewall_port_direction: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.target_firewall_nodename, + config.firewall_port_name, + config.firewall_port_direction, + "acl", + "add_rule", + config.permission, + config.protocol_name, + str(config.src_ip), + str(config.src_wildcard), + config.src_port, + str(config.dst_ip), + str(config.dst_wildcard), + config.dst_port, + config.position, + ] + + +class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="firewall_acl_remove_rule"): + """Action which removes a rule from a firewall port's ACL.""" + + config: "FirewallACLRemoveRuleAction.ConfigSchema" + + class ConfigSchema(ACLRemoveRuleAbstractAction.ConfigSchema): + """Configuration schema for FirewallACLRemoveRuleAction.""" + + target_firewall_nodename: str + firewall_port_name: str + firewall_port_direction: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.target_firewall_nodename, + config.firewall_port_name, + config.firewall_port_direction, + "acl", + "remove_rule", + config.position, + ] diff --git a/src/primaite/game/agent/actions/application.py b/src/primaite/game/agent/actions/application.py new file mode 100644 index 00000000..f6ce0624 --- /dev/null +++ b/src/primaite/game/agent/actions/application.py @@ -0,0 +1,137 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import ABC +from typing import ClassVar + +from primaite.game.agent.actions.abstract import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "NodeApplicationExecuteAction", + "NodeApplicationScanAction", + "NodeApplicationCloseAction", + "NodeApplicationFixAction", + "NodeApplicationInstallAction", + "NodeApplicationRemoveAction", +) + + +class NodeApplicationAbstractAction(AbstractAction, ABC): + """ + Base class for application actions. + + Any action which applies to an application and uses node_name and application_name as its only two parameters can + inherit from this base class. + """ + + config: "NodeApplicationAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base Configuration schema for Node Application actions.""" + + node_name: str + application_name: str + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.node_name, + "application", + config.application_name, + config.verb, + ] + + +class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="node_application_execute"): + """Action which executes an application.""" + + config: "NodeApplicationExecuteAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationExecuteAction.""" + + verb: str = "execute" + + +class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_application_scan"): + """Action which scans an application.""" + + config: "NodeApplicationScanAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationScanAction.""" + + verb: str = "scan" + + +class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node_application_close"): + """Action which closes an application.""" + + config: "NodeApplicationCloseAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationCloseAction.""" + + verb: str = "close" + + +class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_application_fix"): + """Action which fixes an application.""" + + config: "NodeApplicationFixAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationFixAction.""" + + verb: str = "fix" + + +class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="node_application_install"): + """Action which installs an application.""" + + config: "NodeApplicationInstallAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationInstallAction.""" + + verb: str = "install" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.node_name, + "software_manager", + "application", + config.verb, + config.application_name, + ] + + +class NodeApplicationRemoveAction(NodeApplicationAbstractAction, identifier="node_application_remove"): + """Action which removes/uninstalls an application.""" + + config: "NodeApplicationRemoveAction.ConfigSchema" + + class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema): + """Configuration schema for NodeApplicationRemoveAction.""" + + verb: str = "uninstall" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.node_name, + "software_manager", + "application", + config.verb, + config.application_name, + ] diff --git a/src/primaite/game/agent/actions/file.py b/src/primaite/game/agent/actions/file.py new file mode 100644 index 00000000..ed666773 --- /dev/null +++ b/src/primaite/game/agent/actions/file.py @@ -0,0 +1,189 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import ABC +from typing import ClassVar + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "NodeFileCreateAction", + "NodeFileScanAction", + "NodeFileDeleteAction", + "NodeFileRestoreAction", + "NodeFileCorruptAction", + "NodeFileAccessAction", + "NodeFileCheckhashAction", + "NodeFileRepairAction", +) + + +class NodeFileAbstractAction(AbstractAction, ABC): + """Abstract base class for file actions. + + Any action which applies to a file and uses node_name, folder_name, and file_name as its + only three parameters can inherit from this base class. + """ + + config: "NodeFileAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration Schema for NodeFileAbstractAction.""" + + node_name: str + folder_name: str + file_name: str + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None or config.file_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + "folder", + config.folder_name, + "file", + config.file_name, + config.verb, + ] + + +class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"): + """Action which creates a new file in a given folder.""" + + config: "NodeFileCreateAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileCreateAction.""" + + verb: ClassVar[str] = "create" + force: bool = False + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None or config.file_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + config.verb, + "file", + config.folder_name, + config.file_name, + config.verb, + ] + + +class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"): + """Action which scans a file.""" + + config: "NodeFileScanAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileScanAction.""" + + verb: ClassVar[str] = "scan" + + +class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"): + """Action which deletes a file.""" + + config: "NodeFileDeleteAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileDeleteAction.""" + + verb: ClassVar[str] = "delete" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None or config.file_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + config.verb, + "file", + config.folder_name, + config.file_name, + ] + + +class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restore"): + """Action which restores a file.""" + + config: "NodeFileRestoreAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileRestoreAction.""" + + verb: ClassVar[str] = "restore" + + +class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrupt"): + """Action which corrupts a file.""" + + config: "NodeFileCorruptAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileCorruptAction.""" + + verb: ClassVar[str] = "corrupt" + + +class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"): + """Action which increases a file's access count.""" + + config: "NodeFileAccessAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileAccessAction.""" + + verb: ClassVar[str] = "access" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None or config.file_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + config.verb, + config.folder_name, + config.file_name, + ] + + +class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_checkhash"): + """Action which checks the hash of a file.""" + + config: "NodeFileCheckhashAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration schema for NodeFileCheckhashAction.""" + + verb: ClassVar[str] = "checkhash" + + +class NodeFileRepairAction(NodeFileAbstractAction, identifier="node_file_repair"): + """Action which repairs a file.""" + + config: "NodeFileRepairAction.ConfigSchema" + + class ConfigSchema(NodeFileAbstractAction.ConfigSchema): + """Configuration Schema for NodeFileRepairAction.""" + + verb: ClassVar[str] = "repair" diff --git a/src/primaite/game/agent/actions/folder.py b/src/primaite/game/agent/actions/folder.py new file mode 100644 index 00000000..3e1136ac --- /dev/null +++ b/src/primaite/game/agent/actions/folder.py @@ -0,0 +1,117 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import ABC +from typing import ClassVar + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "NodeFolderScanAction", + "NodeFolderCheckhashAction", + "NodeFolderRepairAction", + "NodeFolderRestoreAction", + "NodeFolderCreateAction", +) + + +class NodeFolderAbstractAction(AbstractAction, ABC): + """ + Base class for folder actions. + + Any action which applies to a folder and uses node_name and folder_name as its only two parameters can inherit from + this base class. + """ + + config: "NodeFolderAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base configuration schema for NodeFolder actions.""" + + node_name: str + folder_name: str + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + "folder", + config.folder_name, + config.verb, + ] + + +class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"): + """Action which scans a folder.""" + + config: "NodeFolderScanAction.ConfigSchema" + + class ConfigSchema(NodeFolderAbstractAction.ConfigSchema): + """Configuration schema for NodeFolderScanAction.""" + + verb: ClassVar[str] = "scan" + + +class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folder_checkhash"): + """Action which checks the hash of a folder.""" + + config: "NodeFolderCheckhashAction.ConfigSchema" + + class ConfigSchema(NodeFolderAbstractAction.ConfigSchema): + """Configuration schema for NodeFolderCheckhashAction.""" + + verb: ClassVar[str] = "checkhash" + + +class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_repair"): + """Action which repairs a folder.""" + + config: "NodeFolderRepairAction.ConfigSchema" + + class ConfigSchema(NodeFolderAbstractAction.ConfigSchema): + """Configuration schema for NodeFolderRepairAction.""" + + verb: ClassVar[str] = "repair" + + +class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_restore"): + """Action which restores a folder.""" + + config: "NodeFolderRestoreAction.ConfigSchema" + + class ConfigSchema(NodeFolderAbstractAction.ConfigSchema): + """Configuration schema for NodeFolderRestoreAction.""" + + verb: ClassVar[str] = "restore" + + +class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_create"): + """Action which creates a new folder.""" + + config: "NodeFolderCreateAction.ConfigSchema" + + class ConfigSchema(NodeFolderAbstractAction.ConfigSchema): + """Configuration schema for NodeFolderCreateAction.""" + + verb: ClassVar[str] = "create" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + config.verb, + "folder", + config.folder_name, + ] diff --git a/src/primaite/game/agent/actions/host_nic.py b/src/primaite/game/agent/actions/host_nic.py new file mode 100644 index 00000000..0ca816f3 --- /dev/null +++ b/src/primaite/game/agent/actions/host_nic.py @@ -0,0 +1,62 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import ABC +from typing import ClassVar + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ("HostNICEnableAction", "HostNICDisableAction") + + +class HostNICAbstractAction(AbstractAction, ABC): + """ + Abstract base class for NIC actions. + + Any action which applies to a NIC and uses node_name and nic_num as its only two parameters can inherit from this + base class. + """ + + config: "HostNICAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base Configuration schema for HostNIC actions.""" + + node_name: str + nic_num: int + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.nic_num is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "network_interface", + config.nic_num, + config.verb, + ] + + +class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"): + """Action which enables a NIC.""" + + config: "HostNICEnableAction.ConfigSchema" + + class ConfigSchema(HostNICAbstractAction.ConfigSchema): + """Configuration schema for HostNICEnableAction.""" + + verb: ClassVar[str] = "enable" + + +class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"): + """Action which disables a NIC.""" + + config: "HostNICDisableAction.ConfigSchema" + + class ConfigSchema(HostNICAbstractAction.ConfigSchema): + """Configuration schema for HostNICDisableAction.""" + + verb: ClassVar[str] = "disable" diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py new file mode 100644 index 00000000..3e5b21b1 --- /dev/null +++ b/src/primaite/game/agent/actions/manager.py @@ -0,0 +1,108 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +"""yaml example. + +agents: + - name: agent_1 + action_space: + actions: + - do_nothing + - node_service_start + - node_service_stop + action_map: +""" + +from __future__ import annotations + +from typing import Dict, Tuple + +from gymnasium import spaces +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from primaite.game.agent.actions.abstract import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ("DoNothingAction", "ActionManager") + + +class DoNothingAction(AbstractAction, identifier="do_nothing"): + """Do Nothing Action.""" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration Schema for do_nothingAction.""" + + type: str = "do_nothing" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return ["do_nothing"] + + +class _ActionMapItem(BaseModel): + model_config = ConfigDict(extra="forbid") + + action: str + options: Dict + + +class ActionManager(BaseModel): + """Class which manages the action space for an agent.""" + + class ConfigSchema(BaseModel): + """Config Schema for ActionManager.""" + + model_config = ConfigDict(extra="forbid") + action_map: Dict[int, _ActionMapItem] = {} + """Mapping between integer action choices and CAOS actions.""" + + @field_validator("action_map", mode="after") + def consecutive_action_nums(cls, v: Dict) -> Dict: + """Make sure all numbers between 0 and N are represented as dict keys in action map.""" + assert all([i in v.keys() for i in range(len(v))]) + return v + + config: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema()) + + action_map: Dict[int, Tuple[str, Dict]] = {} + """Init as empty, populate after model validation.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.action_map = {n: (v.action, v.options) for n, v in self.config.action_map.items()} + + def get_action(self, action: int) -> Tuple[str, Dict]: + """ + Produce action in CAOS format. + + The agent chooses an action (as an integer), this is converted into an action in CAOS format + The CAOS format is basically an action identifier, followed by parameters stored in a dictionary. + """ + act_identifier, act_options = self.action_map[action] + return act_identifier, act_options + + def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat: + """Take action in CAOS format and use the execution definition to change it into PrimAITE request format.""" + act_class = AbstractAction._registry[action_identifier] + config = act_class.ConfigSchema(**action_options) + return act_class.form_request(config=config) + + @property + def space(self) -> spaces.Space: + """Return the gymnasium action space for this agent.""" + return spaces.Discrete(len(self.action_map)) + + @classmethod + def from_config(cls, cfg: Dict) -> "ActionManager": + """ + Construct an ActionManager from a config dictionary. + + The action space config supports must contain the following key: + ``action_map`` - List of actions available to the agent, formatted as a dictionary where the key is the + action number between 0 - N, and the value is the CAOS-formatted action. + + :param cfg: The action space config. + :type cfg: Dict + :return: The constructed ActionManager. + :rtype: ActionManager + """ + return cls(**cfg.get("options", {}), act_map=cfg.get("action_map")) diff --git a/src/primaite/game/agent/actions/network.py b/src/primaite/game/agent/actions/network.py new file mode 100644 index 00000000..d244fb74 --- /dev/null +++ b/src/primaite/game/agent/actions/network.py @@ -0,0 +1,57 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + +from typing import ClassVar + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction") + + +class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstract"): + """Base class for Network port actions.""" + + config: "NetworkPortAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base configuration schema for NetworkPort actions.""" + + target_nodename: str + port_num: int + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.target_nodename is None or config.port_num is None: + return ["do_nothing"] + return [ + "network", + "node", + config.target_nodename, + "network_interface", + config.port_num, + config.verb, + ] + + +class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"): + """Action which enables are port on a router or a firewall.""" + + config: "NetworkPortEnableAction.ConfigSchema" + + class ConfigSchema(NetworkPortAbstractAction.ConfigSchema): + """Configuration schema for NetworkPortEnableAction.""" + + verb: ClassVar[str] = "enable" + + +class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"): + """Action which disables are port on a router or a firewall.""" + + config: "NetworkPortDisableAction.ConfigSchema" + + class ConfigSchema(NetworkPortAbstractAction.ConfigSchema): + """Configuration schema for NetworkPortDisableAction.""" + + verb: ClassVar[str] = "disable" diff --git a/src/primaite/game/agent/actions/node.py b/src/primaite/game/agent/actions/node.py new file mode 100644 index 00000000..fbab18f0 --- /dev/null +++ b/src/primaite/game/agent/actions/node.py @@ -0,0 +1,186 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import abstractmethod +from typing import ClassVar, List, Optional, Union + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat +from primaite.utils.validation.ip_protocol import IPProtocol +from primaite.utils.validation.port import Port + +__all__ = ( + "NodeOSScanAction", + "NodeShutdownAction", + "NodeStartupAction", + "NodeResetAction", + "NodeNMAPPingScanAction", + "NodeNMAPPortScanAction", + "NodeNetworkServiceReconAction", +) + + +class NodeAbstractAction(AbstractAction, identifier="node_abstract"): + """ + Abstract base class for node actions. + + Any action which applies to a node and uses node_name as its only parameter can inherit from this base class. + """ + + config: "NodeAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base Configuration schema for Node actions.""" + + node_name: str + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + print(config) + return ["network", "node", config.node_name, config.verb] + + +class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"): + """Action which scans a node's OS.""" + + config: "NodeOSScanAction.ConfigSchema" + + class ConfigSchema(NodeAbstractAction.ConfigSchema): + """Configuration schema for NodeOSScanAction.""" + + verb: ClassVar[str] = "scan" + + +class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"): + """Action which shuts down a node.""" + + config: "NodeShutdownAction.ConfigSchema" + + class ConfigSchema(NodeAbstractAction.ConfigSchema): + """Configuration schema for NodeShutdownAction.""" + + verb: ClassVar[str] = "shutdown" + + +class NodeStartupAction(NodeAbstractAction, identifier="node_startup"): + """Action which starts up a node.""" + + config: "NodeStartupAction.ConfigSchema" + + class ConfigSchema(NodeAbstractAction.ConfigSchema): + """Configuration schema for NodeStartupAction.""" + + verb: ClassVar[str] = "startup" + + +class NodeResetAction(NodeAbstractAction, identifier="node_reset"): + """Action which resets a node.""" + + config: "NodeResetAction.ConfigSchema" + + class ConfigSchema(NodeAbstractAction.ConfigSchema): + """Configuration schema for NodeResetAction.""" + + verb: ClassVar[str] = "reset" + + +class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_action"): + """Base class for NodeNMAP actions.""" + + config: "NodeNMAPAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base Configuration Schema for NodeNMAP actions.""" + + target_ip_address: Union[str, List[str]] + show: bool = False + source_node: str + + @classmethod + @abstractmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + # NMAP action requests don't share a common format for their requests + # This is just a placeholder to ensure the method is defined. + pass + + +class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_scan"): + """Action which performs an NMAP ping scan.""" + + config: "NodeNMAPPingScanAction.ConfigSchema" + + @classmethod + def form_request(cls, config: "NodeNMAPPingScanAction.ConfigSchema") -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.source_node, + "application", + "NMAP", + "ping_scan", + {"target_ip_address": config.target_ip_address, "show": config.show}, + ] + + +class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_scan"): + """Action which performs an NMAP port scan.""" + + config: "NodeNMAPPortScanAction.ConfigSchema" + + class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema): + """Configuration Schema for NodeNMAPPortScanAction.""" + + source_node: str + target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None + target_port: Optional[Union[Port, List[Port]]] = None + show: Optional[bool] = (False,) + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.source_node, + "application", + "NMAP", + "port_scan", + { + "target_ip_address": config.target_ip_address, + "target_port": config.target_port, + "target_protocol": config.target_protocol, + "show": config.show, + }, + ] + + +class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_network_service_recon"): + """Action which performs an NMAP network service recon (ping scan followed by port scan).""" + + config: "NodeNetworkServiceReconAction.ConfigSchema" + + class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema): + """Configuration schema for NodeNetworkServiceReconAction.""" + + target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None + target_port: Optional[Union[Port, List[Port]]] = None + show: Optional[bool] = (False,) + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.source_node, + "application", + "NMAP", + "network_service_recon", + { + "target_ip_address": config.target_ip_address, + "target_port": config.target_port, + "target_protocol": config.target_protocol, + "show": config.show, + }, + ] diff --git a/src/primaite/game/agent/actions/service.py b/src/primaite/game/agent/actions/service.py new file mode 100644 index 00000000..4a483f28 --- /dev/null +++ b/src/primaite/game/agent/actions/service.py @@ -0,0 +1,135 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from typing import ClassVar + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "NodeServiceScanAction", + "NodeServiceStopAction", + "NodeServiceStartAction", + "NodeServicePauseAction", + "NodeServiceResumeAction", + "NodeServiceRestartAction", + "NodeServiceDisableAction", + "NodeServiceEnableAction", + "NodeServiceFixAction", +) + + +class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstract"): + """Abstract Action for Node Service related actions. + + Any actions which use node_name and service_name can inherit from this class. + """ + + config: "NodeServiceAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + node_name: str + service_name: str + verb: ClassVar[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return ["network", "node", config.node_name, "service", config.service_name, config.verb] + + +class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_scan"): + """Action which scans a service.""" + + config: "NodeServiceScanAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceScanAction.""" + + verb: ClassVar[str] = "scan" + + +class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_stop"): + """Action which stops a service.""" + + config: "NodeServiceStopAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceStopAction.""" + + verb: ClassVar[str] = "stop" + + +class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service_start"): + """Action which starts a service.""" + + config: "NodeServiceStartAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceStartAction.""" + + verb: ClassVar[str] = "start" + + +class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service_pause"): + """Action which pauses a service.""" + + config: "NodeServicePauseAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServicePauseAction.""" + + verb: ClassVar[str] = "pause" + + +class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_service_resume"): + """Action which resumes a service.""" + + config: "NodeServiceResumeAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceResumeAction.""" + + verb: ClassVar[str] = "resume" + + +class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_service_restart"): + """Action which restarts a service.""" + + config: "NodeServiceRestartAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceRestartAction.""" + + verb: ClassVar[str] = "restart" + + +class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_service_disable"): + """Action which disables a service.""" + + config: "NodeServiceDisableAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceDisableAction.""" + + verb: ClassVar[str] = "disable" + + +class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_service_enable"): + """Action which enables a service.""" + + config: "NodeServiceEnableAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceEnableAction.""" + + verb: ClassVar[str] = "enable" + + +class NodeServiceFixAction(NodeServiceAbstractAction, identifier="node_service_fix"): + """Action which fixes a service.""" + + config: "NodeServiceFixAction.ConfigSchema" + + class ConfigSchema(NodeServiceAbstractAction.ConfigSchema): + """Configuration Schema for NodeServiceFixAction.""" + + verb: ClassVar[str] = "fix" diff --git a/src/primaite/game/agent/actions/session.py b/src/primaite/game/agent/actions/session.py new file mode 100644 index 00000000..9720d371 --- /dev/null +++ b/src/primaite/game/agent/actions/session.py @@ -0,0 +1,108 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from abc import abstractmethod + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "NodeSessionsRemoteLoginAction", + "NodeSessionsRemoteLogoutAction", + "NodeAccountChangePasswordAction", +) + + +class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"): + """Base class for NodeSession actions.""" + + config: "NodeSessionAbstractAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base configuration schema for NodeSessionAbstractActions.""" + + node_name: str + remote_ip: str + + @classmethod + @abstractmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """ + Abstract method for request forming. + + Should return the action formatted as a request which can be ingested by the PrimAITE simulation. + """ + pass + + +class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_session_remote_login"): + """Action which performs a remote session login.""" + + config: "NodeSessionsRemoteLoginAction.ConfigSchema" + + class ConfigSchema(NodeSessionAbstractAction.ConfigSchema): + """Configuration schema for NodeSessionsRemoteLoginAction.""" + + username: str + password: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.remote_ip is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "service", + "Terminal", + "node_session_remote_login", + config.username, + config.password, + config.remote_ip, + ] + + +class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node_session_remote_logoff"): + """Action which performs a remote session logout.""" + + config: "NodeSessionsRemoteLogoutAction.ConfigSchema" + + class ConfigSchema(NodeSessionAbstractAction.ConfigSchema): + """Configuration schema for NodeSessionsRemoteLogoutAction.""" + + verb: str = "remote_logoff" + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.remote_ip is None: + return ["do_nothing"] + return ["network", "node", config.node_name, "service", "Terminal", config.verb, config.remote_ip] + + +class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"): + """Action which changes the password for a user.""" + + config: "NodeAccountChangePasswordAction.ConfigSchema" + + class ConfigSchema(NodeSessionAbstractAction.ConfigSchema): + """Configuration schema for NodeAccountsChangePasswordAction.""" + + username: str + current_password: str + new_password: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.node_name, + "service", + "UserManager", + "change_password", + config.username, + config.current_password, + config.new_password, + ] diff --git a/src/primaite/game/agent/actions/software.py b/src/primaite/game/agent/actions/software.py new file mode 100644 index 00000000..e0d602ed --- /dev/null +++ b/src/primaite/game/agent/actions/software.py @@ -0,0 +1,241 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + +from typing import List, Optional, Union + +from pydantic import ConfigDict, Field + +from primaite.game.agent.actions.manager import AbstractAction +from primaite.interface.request import RequestFormat + +__all__ = ( + "ConfigureRansomwareScriptAction", + "ConfigureDoSBotAction", + "ConfigureC2BeaconAction", + "NodeSendRemoteCommandAction", + "TerminalC2ServerAction", + "RansomwareLaunchC2ServerAction", + "ExfiltrationC2ServerAction", + "ConfigureDatabaseClientAction", +) + + +class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_ransomware_script"): + """Action which sets config parameters for a ransomware script on a node.""" + + config: "ConfigureRansomwareScriptAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for ConfigureRansomwareScriptAction.""" + + node_name: str + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + payload: Optional[str] = None + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + if config.node_name is None: + return ["do_nothing"] + data = dict( + server_ip_address=config.server_ip_address, + server_password=config.server_password, + payload=config.payload, + ) + return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data] + + +class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identifier="c2_server_ransomware_configure"): + """Action which causes a C2 server to send a command to set options on a ransomware script remotely.""" + + @classmethod + def form_request(cls, config: ConfigureRansomwareScriptAction.ConfigSchema) -> RequestFormat: + data = dict( + server_ip_address=config.server_ip_address, server_password=config.server_password, payload=config.payload + ) + return ["network", "node", config.node_name, "application", "C2Server", "ransomware_configure", data] + + +class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): + """Action which sets config parameters for a DoS bot on a node.""" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Schema for options that can be passed to this action.""" + + model_config = ConfigDict(extra="forbid") + node_name: str + target_ip_address: Optional[str] = None + target_port: Optional[str] = None + payload: Optional[str] = None + repeat: Optional[bool] = None + port_scan_p_of_success: Optional[float] = None + dos_intensity: Optional[float] = None + max_sessions: Optional[int] = None + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + data = dict( + target_ip_address=config.target_ip_address, + target_port=config.target_port, + payload=config.payload, + repeat=config.repeat, + port_scan_p_of_success=config.port_scan_p_of_success, + dos_intensity=config.dos_intensity, + max_sessions=config.max_sessions, + ) + data = {k: v for k, v in data.items() if v is not None} + return ["network", "node", config.node_name, "application", "DoSBot", "configure", data] + + +class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"): + """Action which configures a C2 Beacon based on the parameters given.""" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for ConfigureC2BeaconAction.""" + + node_name: str + c2_server_ip_address: str + keep_alive_frequency: int = Field(default=5, ge=1) + masquerade_protocol: str = Field(default="TCP") + masquerade_port: str = Field(default="HTTP") + + @classmethod + def form_request(self, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + data = dict( + c2_server_ip_address=config.c2_server_ip_address, + keep_alive_frequency=config.keep_alive_frequency, + masquerade_protocol=config.masquerade_protocol, + masquerade_port=config.masquerade_port, + ) + return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data] + + +class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"): + """Action which sends a terminal command to a remote node via SSH.""" + + config: "NodeSendRemoteCommandAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for NodeSendRemoteCommandAction.""" + + node_name: str + remote_ip: str + command: RequestFormat + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [ + "network", + "node", + config.node_name, + "service", + "Terminal", + "send_remote_command", + config.remote_ip, + {"command": config.command}, + ] + + +class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_command"): + """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" + + config: "TerminalC2ServerAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Schema for options that can be passed to this action.""" + + node_name: str + commands: Union[List[RequestFormat], RequestFormat] + ip_address: Optional[str] + username: Optional[str] + password: Optional[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + if config.node_name is None: + return ["do_nothing"] + + command_model = { + "commands": config.commands, + "ip_address": config.ip_address, + "username": config.username, + "password": config.password, + } + return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model] + + +class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ransomware_launch"): + """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" + + config: "RansomwareLaunchC2ServerAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for RansomwareLaunchC2ServerAction.""" + + node_name: str + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + if config.node_name is None: + return ["do_nothing"] + # This action currently doesn't require any further configuration options. + return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"] + + +class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfiltrate"): + """Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server.""" + + config: "ExfiltrationC2ServerAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Schema for options that can be passed to this action.""" + + node_name: str + username: Optional[str] + password: Optional[str] + target_ip_address: str + target_file_name: str + target_folder_name: str + exfiltration_folder_name: Optional[str] + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + if config.node_name is None: + return ["do_nothing"] + + command_model = { + "target_file_name": config.target_file_name, + "target_folder_name": config.target_folder_name, + "exfiltration_folder_name": config.exfiltration_folder_name, + "target_ip_address": config.target_ip_address, + "username": config.username, + "password": config.password, + } + return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model] + + +class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"): + """Action which sets config parameters for a database client on a node.""" + + config: "ConfigureDatabaseClientAction.ConfigSchema" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Schema for options that can be passed to this action.""" + + node_name: str + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + if config.node_name is None: + return ["do_nothing"] + data = {"server_ip_address": config.server_ip_address, "server_password": config.server_password} + return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", data] diff --git a/src/primaite/game/agent/agent_log.py b/src/primaite/game/agent/agent_log.py index 59fb4702..5d9dc848 100644 --- a/src/primaite/game/agent/agent_log.py +++ b/src/primaite/game/agent/agent_log.py @@ -1,6 +1,7 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK import logging from pathlib import Path +from typing import Optional from prettytable import MARKDOWN, PrettyTable @@ -20,20 +21,22 @@ class _NotJSONFilter(logging.Filter): class AgentLog: """ - A Agent Log class is a simple logger dedicated to managing and writing logging updates and information for an agent. + An Agent Log class is a simple logger dedicated to managing and writing updates and information for an agent. - Each log message is written to a file located at: /agent_name/agent_name.log + Each log message is written to a file located at: + /agent_name/agent_name.log """ - def __init__(self, agent_name: str): + def __init__(self, agent_name: Optional[str]): """ Constructs a Agent Log instance for a given hostname. - :param hostname: The hostname associated with the system logs being recorded. + :param agent_name: The agent_name associated with the system logs being recorded. """ - self.agent_name = agent_name - self.current_episode: int = 1 + super().__init__() + self.agent_name = agent_name if agent_name else "unnamed_agent" self.current_timestep: int = 0 + self.current_episode: int = 1 self.setup_logger() @property @@ -90,7 +93,7 @@ class AgentLog: def _write_to_terminal(self, msg: str, level: str, to_terminal: bool = False): if to_terminal or SIM_OUTPUT.write_agent_log_to_terminal: - print(f"{self.agent_name}: ({ self.timestep}) ({level}) {msg}") + print(f"{self.agent_name}: ({self.timestep}) ({level}) {msg}") def debug(self, msg: str, to_terminal: bool = False): """ diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 4acc9108..aac898e1 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -1,10 +1,12 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK """Interface for agents.""" +from __future__ import annotations + from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING +from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, TYPE_CHECKING from gymnasium.core import ActType, ObsType -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, ConfigDict, Field from primaite.game.agent.actions import ActionManager from primaite.game.agent.agent_log import AgentLog @@ -15,6 +17,8 @@ from primaite.interface.request import RequestFormat, RequestResponse if TYPE_CHECKING: pass +__all__ = ("AgentHistoryItem", "AbstractAgent", "AbstractScriptedAgent", "ProxyAgent") + class AgentHistoryItem(BaseModel): """One entry of an agent's action log - what the agent did and how the simulator responded in 1 step.""" @@ -39,89 +43,56 @@ class AgentHistoryItem(BaseModel): reward_info: Dict[str, Any] = {} -class AgentStartSettings(BaseModel): - """Configuration values for when an agent starts performing actions.""" - - start_step: int = 5 - "The timestep at which an agent begins performing it's actions" - frequency: int = 5 - "The number of timesteps to wait between performing actions" - variance: int = 0 - "The amount the frequency can randomly change to" - - @model_validator(mode="after") - def check_variance_lt_frequency(self) -> "AgentStartSettings": - """ - Make sure variance is equal to or lower than frequency. - - This is because the calculation for the next execution time is now + (frequency +- variance). If variance were - greater than frequency, sometimes the bracketed term would be negative and the attack would never happen again. - """ - if self.variance > self.frequency: - raise ValueError( - f"Agent start settings error: variance must be lower than frequency " - f"{self.variance=}, {self.frequency=}" - ) - return self - - -class AgentSettings(BaseModel): - """Settings for configuring the operation of an agent.""" - - start_settings: Optional[AgentStartSettings] = None - "Configuration for when an agent begins performing it's actions" - flatten_obs: bool = True - "Whether to flatten the observation space before passing it to the agent. True by default." - action_masking: bool = False - "Whether to return action masks at each step." - - @classmethod - def from_config(cls, config: Optional[Dict]) -> "AgentSettings": - """Construct agent settings from a config dictionary. - - :param config: A dict of options for the agent settings. - :type config: Dict - :return: The agent settings. - :rtype: AgentSettings - """ - if config is None: - return cls() - - return cls(**config) - - -class AbstractAgent(ABC): +class AbstractAgent(BaseModel, ABC): """Base class for scripted and RL agents.""" - def __init__( - self, - agent_name: Optional[str], - action_space: Optional[ActionManager], - observation_space: Optional[ObservationManager], - reward_function: Optional[RewardFunction], - agent_settings: Optional[AgentSettings] = None, - ) -> None: - """ - Initialize an agent. + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes. - :type agent_name: Optional[str] - :param action_space: Action space for the agent. - :type action_space: Optional[ActionManager] - :param observation_space: Observation space for the agent. - :type observation_space: Optional[ObservationSpace] - :param reward_function: Reward function for the agent. - :type reward_function: Optional[RewardFunction] - :param agent_settings: Configurable Options for Abstracted Agents - :type agent_settings: Optional[AgentSettings] - """ - self.agent_name: str = agent_name or "unnamed_agent" - self.action_manager: Optional[ActionManager] = action_space - self.observation_manager: Optional[ObservationManager] = observation_space - self.reward_function: Optional[RewardFunction] = reward_function - self.agent_settings = agent_settings or AgentSettings() - self.history: List[AgentHistoryItem] = [] - self.logger = AgentLog(agent_name) + class AgentSettingsSchema(BaseModel, ABC): + """Schema for the 'agent_settings' key.""" + + model_config = ConfigDict(extra="forbid") + + class ConfigSchema(BaseModel, ABC): + """Configuration Schema for AbstractAgents.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + type: str + ref: str = "" + """name of the agent.""" + team: Optional[Literal["BLUE", "GREEN", "RED"]] = None + agent_settings: AbstractAgent.AgentSettingsSchema = Field(default=lambda: AbstractAgent.AgentSettingsSchema()) + action_space: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema()) + observation_space: ObservationManager.ConfigSchema = Field( + default_factory=lambda: ObservationManager.ConfigSchema() + ) + reward_function: RewardFunction.ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema()) + + config: "AbstractAgent.ConfigSchema" = Field(default_factory=lambda: AbstractAgent.ConfigSchema()) + + logger: AgentLog = AgentLog(agent_name="Abstract_Agent") + history: List[AgentHistoryItem] = [] + + action_manager: ActionManager = Field(default_factory=lambda: ActionManager()) + observation_manager: ObservationManager = Field(default_factory=lambda: ObservationManager()) + reward_function: RewardFunction = Field(default_factory=lambda: RewardFunction()) + + _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} + + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if identifier is None: + return + if identifier in cls._registry: + raise ValueError(f"Cannot create a new agent under reserved name {identifier}") + cls._registry[identifier] = cls + + def model_post_init(self, __context: Any) -> None: + """Overwrite the default empty action, observation, and rewards with ones defined through the config.""" + self.action_manager = ActionManager(config=self.config.action_space) + self.observation_manager = ObservationManager(config=self.config.observation_space) + self.reward_function = RewardFunction(config=self.config.reward_function) + return super().model_post_init(__context) def update_observation(self, state: Dict) -> ObsType: """ @@ -159,9 +130,9 @@ class AbstractAgent(ABC): """ # in RL agent, this method will send CAOS observation to RL agent, then receive a int 0-39, # then use a bespoke conversion to take 1-40 int back into CAOS action - return ("DO_NOTHING", {}) + return ("do_nothing", {}) - def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> List[str]: + def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> RequestFormat: # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR """Format action into format expected by the simulator, and apply execution definition if applicable.""" @@ -182,36 +153,47 @@ class AbstractAgent(ABC): """Update the most recent history item with the reward value.""" self.history[-1].reward = self.reward_function.current_reward + @classmethod + def from_config(cls, config: Dict) -> AbstractAgent: + """Grab the relevant agent class and construct an instance from a config dict.""" + agent_type = config["type"] + agent_class = cls._registry[agent_type] + return agent_class(config=config) -class AbstractScriptedAgent(AbstractAgent): + +class AbstractScriptedAgent(AbstractAgent, identifier="AbstractScriptedAgent"): """Base class for actors which generate their own behaviour.""" + config: "AbstractScriptedAgent.ConfigSchema" = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema()) + + class ConfigSchema(AbstractAgent.ConfigSchema): + """Configuration Schema for AbstractScriptedAgents.""" + + type: str = "AbstractScriptedAgent" + @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): +class ProxyAgent(AbstractAgent, identifier="ProxyAgent"): """Agent that sends observations to an RL model and receives actions from that model.""" - def __init__( - self, - agent_name: Optional[str], - action_space: Optional[ActionManager], - observation_space: Optional[ObservationManager], - reward_function: Optional[RewardFunction], - agent_settings: Optional[AgentSettings] = None, - ) -> None: - super().__init__( - agent_name=agent_name, - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - ) - self.most_recent_action: ActType - self.flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False - self.action_masking: bool = agent_settings.action_masking if agent_settings else False + config: "ProxyAgent.ConfigSchema" = Field(default_factory=lambda: ProxyAgent.ConfigSchema()) + most_recent_action: ActType = None + + class AgentSettingsSchema(AbstractAgent.AgentSettingsSchema): + """Schema for the `agent_settings` part of the agent config.""" + + flatten_obs: bool = False + action_masking: bool = False + + class ConfigSchema(AbstractAgent.ConfigSchema): + """Configuration Schema for Proxy Agent.""" + + type: str = "Proxy_Agent" + agent_settings: ProxyAgent.AgentSettingsSchema = Field(default_factory=lambda: ProxyAgent.AgentSettingsSchema()) def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: """ @@ -233,3 +215,8 @@ class ProxyAgent(AbstractAgent): The environment is responsible for calling this method when it receives an action from the agent policy. """ self.most_recent_action = action + + @property + def flatten_obs(self) -> bool: + """Return agent flatten_obs param.""" + return self.config.agent_settings.flatten_obs diff --git a/src/primaite/game/agent/observations/__init__.py b/src/primaite/game/agent/observations/__init__.py index 68435eed..a38095b3 100644 --- a/src/primaite/game/agent/observations/__init__.py +++ b/src/primaite/game/agent/observations/__init__.py @@ -17,5 +17,5 @@ from primaite.game.agent.observations.software_observation import ApplicationObs __all__ = [ "ACLObservation", "FileObservation", "FolderObservation", "FirewallObservation", "HostObservation", "LinksObservation", "NICObservation", "PortObservation", "NodesObservation", "NestedObservation", - "ObservationManager", "ApplicationObservation", "ServiceObservation",] + "ObservationManager", "ApplicationObservation", "ServiceObservation", "RouterObservation", "LinkObservation",] # fmt: on diff --git a/src/primaite/game/agent/observations/acl_observation.py b/src/primaite/game/agent/observations/acl_observation.py index 86a6463a..cb2cb38e 100644 --- a/src/primaite/game/agent/observations/acl_observation.py +++ b/src/primaite/game/agent/observations/acl_observation.py @@ -24,8 +24,8 @@ class ACLObservation(AbstractObservation, identifier="ACL"): """List of IP addresses.""" wildcard_list: Optional[List[str]] = None """List of wildcard strings.""" - port_list: Optional[List[int]] = None - """List of port numbers.""" + port_list: Optional[List[str]] = None + """List of port names.""" protocol_list: Optional[List[str]] = None """List of protocol names.""" num_rules: Optional[int] = None @@ -37,7 +37,7 @@ class ACLObservation(AbstractObservation, identifier="ACL"): num_rules: int, ip_list: List[IPv4Address], wildcard_list: List[str], - port_list: List[int], + port_list: List[str], protocol_list: List[str], ) -> None: """ @@ -51,8 +51,8 @@ class ACLObservation(AbstractObservation, identifier="ACL"): :type ip_list: List[IPv4Address] :param wildcard_list: List of wildcard strings. :type wildcard_list: List[str] - :param port_list: List of port numbers. - :type port_list: List[int] + :param port_list: List of port names. + :type port_list: List[str] :param protocol_list: List of protocol names. :type protocol_list: List[str] """ @@ -60,7 +60,7 @@ class ACLObservation(AbstractObservation, identifier="ACL"): self.num_rules: int = num_rules self.ip_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(ip_list)} self.wildcard_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(wildcard_list)} - self.port_to_id: Dict[int, int] = {p: i + 2 for i, p in enumerate(port_list)} + self.port_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(port_list)} self.protocol_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(protocol_list)} self.default_observation: Dict = { i diff --git a/src/primaite/game/agent/observations/file_system_observations.py b/src/primaite/game/agent/observations/file_system_observations.py index 50ca93fd..784eaa7f 100644 --- a/src/primaite/game/agent/observations/file_system_observations.py +++ b/src/primaite/game/agent/observations/file_system_observations.py @@ -190,6 +190,8 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"): if self.files: self.default_observation["FILES"] = {i + 1: f.default_observation for i, f in enumerate(self.files)} + self.cached_obs: Optional[ObsType] = self.default_observation + def observe(self, state: Dict) -> ObsType: """ Generate observation based on the current state of the simulation. @@ -204,7 +206,10 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"): return self.default_observation if self.file_system_requires_scan: - health_status = folder_state["visible_status"] + if not folder_state["scanned_this_step"]: + health_status = self.cached_obs["health_status"] + else: + health_status = folder_state["visible_status"] else: health_status = folder_state["health_status"] diff --git a/src/primaite/game/agent/observations/firewall_observation.py b/src/primaite/game/agent/observations/firewall_observation.py index a89ddfc5..44541f24 100644 --- a/src/primaite/game/agent/observations/firewall_observation.py +++ b/src/primaite/game/agent/observations/firewall_observation.py @@ -27,13 +27,13 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): """List of IP addresses for encoding ACLs.""" wildcard_list: Optional[List[str]] = None """List of IP wildcards for encoding ACLs.""" - port_list: Optional[List[int]] = None + port_list: Optional[List[str]] = None """List of ports for encoding ACLs.""" protocol_list: Optional[List[str]] = None """List of protocols for encoding ACLs.""" num_rules: Optional[int] = None """Number of rules ACL rules to show.""" - include_users: Optional[bool] = True + include_users: Optional[bool] = None """If True, report user session information.""" def __init__( @@ -41,7 +41,7 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): where: WhereType, ip_list: List[str], wildcard_list: List[str], - port_list: List[int], + port_list: List[str], protocol_list: List[str], num_rules: int, include_users: bool, @@ -56,8 +56,8 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): :type ip_list: List[str] :param wildcard_list: List of wildcard rules. :type wildcard_list: List[str] - :param port_list: List of port numbers. - :type port_list: List[int] + :param port_list: List of port names. + :type port_list: List[str] :param protocol_list: List of protocol types. :type protocol_list: List[str] :param num_rules: Number of rules configured in the firewall. @@ -72,7 +72,6 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): self.ports: List[PortObservation] = [ PortObservation(where=self.where + ["NICs", port_num]) for port_num in (1, 2, 3) ] - # TODO: check what the port nums are for firewall. self.internal_inbound_acl = ACLObservation( where=self.where + ["internal_inbound_acl", "acl"], @@ -140,6 +139,8 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): }, }, } + if self.include_users: + self.default_observation["users"] = {"local_login": 0, "remote_sessions": 0} def observe(self, state: Dict) -> ObsType: """ @@ -153,29 +154,35 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): firewall_state = access_from_nested_dict(state, self.where) if firewall_state is NOT_PRESENT_IN_STATE: return self.default_observation - obs = { - "PORTS": {i + 1: p.observe(state) for i, p in enumerate(self.ports)}, - "ACL": { - "INTERNAL": { - "INBOUND": self.internal_inbound_acl.observe(state), - "OUTBOUND": self.internal_outbound_acl.observe(state), + + is_on = firewall_state["operating_state"] == 1 + if not is_on: + obs = {**self.default_observation} + + else: + obs = { + "PORTS": {i + 1: p.observe(state) for i, p in enumerate(self.ports)}, + "ACL": { + "INTERNAL": { + "INBOUND": self.internal_inbound_acl.observe(state), + "OUTBOUND": self.internal_outbound_acl.observe(state), + }, + "DMZ": { + "INBOUND": self.dmz_inbound_acl.observe(state), + "OUTBOUND": self.dmz_outbound_acl.observe(state), + }, + "EXTERNAL": { + "INBOUND": self.external_inbound_acl.observe(state), + "OUTBOUND": self.external_outbound_acl.observe(state), + }, }, - "DMZ": { - "INBOUND": self.dmz_inbound_acl.observe(state), - "OUTBOUND": self.dmz_outbound_acl.observe(state), - }, - "EXTERNAL": { - "INBOUND": self.external_inbound_acl.observe(state), - "OUTBOUND": self.external_outbound_acl.observe(state), - }, - }, - } - if self.include_users: - sess = firewall_state["services"]["UserSessionManager"] - obs["users"] = { - "local_login": 1 if sess["current_local_user"] else 0, - "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), } + if self.include_users: + sess = firewall_state["services"]["UserSessionManager"] + obs["users"] = { + "local_login": 1 if sess["current_local_user"] else 0, + "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), + } return obs @property @@ -186,34 +193,36 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"): :return: Gymnasium space representing the observation space for firewall status. :rtype: spaces.Space """ - space = spaces.Dict( - { - "PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}), - "ACL": spaces.Dict( - { - "INTERNAL": spaces.Dict( - { - "INBOUND": self.internal_inbound_acl.space, - "OUTBOUND": self.internal_outbound_acl.space, - } - ), - "DMZ": spaces.Dict( - { - "INBOUND": self.dmz_inbound_acl.space, - "OUTBOUND": self.dmz_outbound_acl.space, - } - ), - "EXTERNAL": spaces.Dict( - { - "INBOUND": self.external_inbound_acl.space, - "OUTBOUND": self.external_outbound_acl.space, - } - ), - } - ), - } - ) - return space + shape = { + "PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}), + "ACL": spaces.Dict( + { + "INTERNAL": spaces.Dict( + { + "INBOUND": self.internal_inbound_acl.space, + "OUTBOUND": self.internal_outbound_acl.space, + } + ), + "DMZ": spaces.Dict( + { + "INBOUND": self.dmz_inbound_acl.space, + "OUTBOUND": self.dmz_outbound_acl.space, + } + ), + "EXTERNAL": spaces.Dict( + { + "INBOUND": self.external_inbound_acl.space, + "OUTBOUND": self.external_outbound_acl.space, + } + ), + } + ), + } + if self.include_users: + shape["users"] = spaces.Dict( + {"local_login": spaces.Discrete(2), "remote_sessions": spaces.Discrete(self.max_users + 1)} + ) + return spaces.Dict(shape) @classmethod def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FirewallObservation: diff --git a/src/primaite/game/agent/observations/host_observations.py b/src/primaite/game/agent/observations/host_observations.py index 03e9aca1..e46cc805 100644 --- a/src/primaite/game/agent/observations/host_observations.py +++ b/src/primaite/game/agent/observations/host_observations.py @@ -54,7 +54,7 @@ class HostObservation(AbstractObservation, identifier="HOST"): """ If True, files and folders must be scanned to update the health state. If False, true state is always shown. """ - include_users: Optional[bool] = True + include_users: Optional[bool] = None """If True, report user session information.""" def __init__( @@ -191,25 +191,31 @@ class HostObservation(AbstractObservation, identifier="HOST"): if node_state is NOT_PRESENT_IN_STATE: return self.default_observation - obs = {} + is_on = node_state["operating_state"] == 1 + if not is_on: + obs = {**self.default_observation} + + else: + obs = {} + if self.services: + obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)} + if self.applications: + obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)} + if self.folders: + obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)} + if self.nics: + obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)} + if self.include_num_access: + obs["num_file_creations"] = node_state["file_system"]["num_file_creations"] + obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"] + if self.include_users: + sess = node_state["services"]["UserSessionManager"] + obs["users"] = { + "local_login": 1 if sess["current_local_user"] else 0, + "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), + } + obs["operating_status"] = node_state["operating_state"] - if self.services: - obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)} - if self.applications: - obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)} - if self.folders: - obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)} - if self.nics: - obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)} - if self.include_num_access: - obs["num_file_creations"] = node_state["file_system"]["num_file_creations"] - obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"] - if self.include_users: - sess = node_state["services"]["UserSessionManager"] - obs["users"] = { - "local_login": 1 if sess["current_local_user"] else 0, - "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), - } return obs @property diff --git a/src/primaite/game/agent/observations/node_observations.py b/src/primaite/game/agent/observations/node_observations.py index 03869367..0c5d11da 100644 --- a/src/primaite/game/agent/observations/node_observations.py +++ b/src/primaite/game/agent/observations/node_observations.py @@ -56,7 +56,7 @@ class NodesObservation(AbstractObservation, identifier="NODES"): """List of IP addresses for encoding ACLs.""" wildcard_list: Optional[List[str]] = None """List of IP wildcards for encoding ACLs.""" - port_list: Optional[List[int]] = None + port_list: Optional[List[str]] = None """List of ports for encoding ACLs.""" protocol_list: Optional[List[str]] = None """List of protocols for encoding ACLs.""" diff --git a/src/primaite/game/agent/observations/observation_manager.py b/src/primaite/game/agent/observations/observation_manager.py index 71a60433..83d4a076 100644 --- a/src/primaite/game/agent/observations/observation_manager.py +++ b/src/primaite/game/agent/observations/observation_manager.py @@ -1,11 +1,12 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from __future__ import annotations +from functools import cached_property from typing import Any, Dict, List, Optional from gymnasium import spaces from gymnasium.core import ObsType -from pydantic import BaseModel, ConfigDict, model_validator, ValidationError +from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validator, ValidationError from primaite.game.agent.observations.observations import AbstractObservation, WhereType @@ -140,7 +141,7 @@ class NullObservation(AbstractObservation, identifier="NONE"): return cls() -class ObservationManager: +class ObservationManager(BaseModel): """ Manage the observations of an Agent. @@ -150,15 +151,66 @@ class ObservationManager: 3. Formatting this information so an agent can use it to make decisions. """ - def __init__(self, obs: AbstractObservation) -> None: - """Initialise observation space. + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - :param observation: Observation object - :type observation: AbstractObservation - """ - self.obs: AbstractObservation = obs - self.current_observation: ObsType - """Cached copy of the observation at the time it was most recently calculated.""" + class ConfigSchema(BaseModel): + """Config Schema for Observation Manager.""" + + model_config = ConfigDict(extra="forbid") + type: str = "NONE" + """Identifier name for the top-level observation.""" + options: AbstractObservation.ConfigSchema = Field( + default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True + ) + """Options to pass into the top-level observation during creation.""" + + @model_validator(mode="before") + @classmethod + def resolve_obs_options_type(cls, data: Any) -> Any: + """ + When constructing the model from a dict, resolve the correct observation class based on `type` field. + + Workaround: The `options` field is statically typed as AbstractObservation. Therefore, it falls over when + passing in data that adheres to a subclass schema rather than the plain AbstractObservation schema. There is + a way to do this properly using discriminated union, but most advice on the internet assumes that the full + list of types between which to discriminate is known ahead-of-time. That is not the case for us, because of + our plugin architecture. + + We may be able to revisit and implement a better solution when needed using the following resources as + research starting points: + https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions + https://github.com/pydantic/pydantic/issues/7366 + https://github.com/pydantic/pydantic/issues/7462 + https://github.com/pydantic/pydantic/pull/7983 + """ + if not isinstance(data, dict): + return data + + # (TODO: duplicate default definition between here and the actual model) + obs_type = data["type"] if "type" in data else "NONE" + obs_class = AbstractObservation._registry[obs_type] + + # if no options are passed in, try to create a default schema. Only works if there are no mandatory fields + if "options" not in data: + data["options"] = obs_class.ConfigSchema() + + # if options passed as a dict, validate against schema + elif isinstance(data["options"], dict): + data["options"] = obs_class.ConfigSchema(**data["options"]) + + return data + + config: ConfigSchema = Field(default_factory=lambda: ObservationManager.ConfigSchema()) + + current_observation: ObsType = 0 + + @computed_field + @cached_property + def obs(self) -> AbstractObservation: + """Create the main observation component for the observation manager from the config.""" + obs_class = AbstractObservation._registry[self.config.type] + obs_instance = obs_class.from_config(config=self.config.options) + return obs_instance def update(self, state: Dict) -> Dict: """ diff --git a/src/primaite/game/agent/observations/observations.py b/src/primaite/game/agent/observations/observations.py index 49b9ab72..89c45b37 100644 --- a/src/primaite/game/agent/observations/observations.py +++ b/src/primaite/game/agent/observations/observations.py @@ -31,7 +31,7 @@ class AbstractObservation(ABC): """Initialise an observation. This method must be overwritten.""" self.default_observation: ObsType - def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: """ Register an observation type. @@ -40,6 +40,8 @@ class AbstractObservation(ABC): :raises ValueError: When attempting to create a component with a name that is already in use. """ super().__init_subclass__(**kwargs) + if identifier is None: + return if identifier in cls._registry: raise ValueError(f"Duplicate observation component type {identifier}") cls._registry[identifier] = cls diff --git a/src/primaite/game/agent/observations/router_observation.py b/src/primaite/game/agent/observations/router_observation.py index ca455f4c..9687d083 100644 --- a/src/primaite/game/agent/observations/router_observation.py +++ b/src/primaite/game/agent/observations/router_observation.py @@ -33,13 +33,13 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"): """List of IP addresses for encoding ACLs.""" wildcard_list: Optional[List[str]] = None """List of IP wildcards for encoding ACLs.""" - port_list: Optional[List[int]] = None + port_list: Optional[List[str]] = None """List of ports for encoding ACLs.""" protocol_list: Optional[List[str]] = None """List of protocols for encoding ACLs.""" num_rules: Optional[int] = None """Number of rules ACL rules to show.""" - include_users: Optional[bool] = True + include_users: Optional[bool] = None """If True, report user session information.""" def __init__( @@ -84,6 +84,8 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"): } if self.ports: self.default_observation["PORTS"] = {i + 1: p.default_observation for i, p in enumerate(self.ports)} + if self.include_users: + self.default_observation["users"] = {"local_login": 0, "remote_sessions": 0} def observe(self, state: Dict) -> ObsType: """ @@ -98,16 +100,21 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"): if router_state is NOT_PRESENT_IN_STATE: return self.default_observation - obs = {} - obs["ACL"] = self.acl.observe(state) - if self.ports: - obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)} - if self.include_users: - sess = router_state["services"]["UserSessionManager"] - obs["users"] = { - "local_login": 1 if sess["current_local_user"] else 0, - "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), - } + is_on = router_state["operating_state"] == 1 + if not is_on: + obs = {**self.default_observation} + + else: + obs = {} + obs["ACL"] = self.acl.observe(state) + if self.ports: + obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)} + if self.include_users: + sess = router_state["services"]["UserSessionManager"] + obs["users"] = { + "local_login": 1 if sess["current_local_user"] else 0, + "remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])), + } return obs @property @@ -121,6 +128,10 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"): shape = {"ACL": self.acl.space} if self.ports: shape["PORTS"] = spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}) + if self.include_users: + shape["users"] = spaces.Dict( + {"local_login": spaces.Discrete(2), "remote_sessions": spaces.Discrete(self.max_users + 1)} + ) return spaces.Dict(shape) @classmethod diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index a4c7c546..80be14ef 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -30,7 +30,7 @@ the structure: from abc import ABC, abstractmethod from typing import Any, Callable, ClassVar, Dict, Iterable, List, Optional, Tuple, Type, TYPE_CHECKING, Union -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing_extensions import Never from primaite import getLogger @@ -48,21 +48,17 @@ class AbstractReward(BaseModel): config: "AbstractReward.ConfigSchema" - # def __init__(self, schema_name, **kwargs): - # super.__init__(self, **kwargs) - # # Create ConfigSchema class - # self.config_class = type(schema_name, (BaseModel, ABC), **kwargs) - # self.config = self.config_class() - class ConfigSchema(BaseModel, ABC): """Config schema for AbstractReward.""" - type: str + type: str = "" _registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {} - def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) + if identifier is None: + return if identifier in cls._registry: raise ValueError(f"Duplicate reward {identifier}") cls._registry[identifier] = cls @@ -381,14 +377,19 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"): class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"): - """Apply a negative reward when taking any action except DONOTHING.""" + """Apply a negative reward when taking any action except do_nothing.""" config: "ActionPenalty.ConfigSchema" class ConfigSchema(AbstractReward.ConfigSchema): - """Config schema for ActionPenalty.""" + """Config schema for ActionPenalty. + + :param action_penalty: Reward to give agents for taking any action except do_nothing + :type action_penalty: float + :param do_nothing_penalty: Reward to give agent for taking the do_nothing action + :type do_nothing_penalty: float + """ - type: str = "ACTION_PENALTY" action_penalty: float = -1.0 do_nothing_penalty: float = 0.0 @@ -402,21 +403,81 @@ class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"): :return: Reward value :rtype: float """ - if last_action_response.action == "DONOTHING": + if last_action_response.action == "do_nothing": return self.config.do_nothing_penalty + else: return self.config.action_penalty -class RewardFunction: +class _SingleComponentConfig(BaseModel): + model_config = ConfigDict(extra="forbid") + type: str + options: AbstractReward.ConfigSchema + weight: float = 1.0 + + @model_validator(mode="before") + @classmethod + def resolve_obs_options_type(cls, data: Any) -> Any: + """ + When constructing the model from a dict, resolve the correct reward class based on `type` field. + + Workaround: The `options` field is statically typed as AbstractReward. Therefore, it falls over when + passing in data that adheres to a subclass schema rather than the plain AbstractReward schema. There is + a way to do this properly using discriminated union, but most advice on the internet assumes that the full + list of types between which to discriminate is known ahead-of-time. That is not the case for us, because of + our plugin architecture. + + We may be able to revisit and implement a better solution when needed using the following resources as + research starting points: + https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions + https://github.com/pydantic/pydantic/issues/7366 + https://github.com/pydantic/pydantic/issues/7462 + https://github.com/pydantic/pydantic/pull/7983 + """ + if not isinstance(data, dict): + return data + + assert "type" in data, ValueError('Reward component definition is missing the "type" key.') + rew_type = data["type"] + rew_class = AbstractReward._registry[rew_type] + + # if no options are passed in, try to create a default schema. Only works if there are no mandatory fields. + if "options" not in data: + data["options"] = rew_class.ConfigSchema() + + # if options are passed as a dict, validate against schema + elif isinstance(data["options"], dict): + data["options"] = rew_class.ConfigSchema(**data["options"]) + + return data + + +class RewardFunction(BaseModel): """Manages the reward function for the agent.""" - def __init__(self): - """Initialise the reward function object.""" - self.reward_components: List[Tuple[AbstractReward, float]] = [] - "attribute reward_components keeps track of reward components and the weights assigned to each." - self.current_reward: float = 0.0 - self.total_reward: float = 0.0 + model_config = ConfigDict(extra="forbid") + + class ConfigSchema(BaseModel): + """Config Schema for RewardFunction.""" + + model_config = ConfigDict(extra="forbid") + + reward_components: Iterable[_SingleComponentConfig] = [] + + config: ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema()) + + reward_components: List[Tuple[AbstractReward, float]] = [] + + current_reward: float = 0.0 + total_reward: float = 0.0 + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + for rew_config in self.config.reward_components: + rew_class = AbstractReward._registry[rew_config.type] + rew_instance = rew_class(config=rew_config.options) + self.register_component(component=rew_instance, weight=rew_config.weight) def register_component(self, component: AbstractReward, weight: float = 1.0) -> None: """Add a reward component to the reward function. diff --git a/src/primaite/game/agent/scripted_agents/__init__.py b/src/primaite/game/agent/scripted_agents/__init__.py index 836b79af..5a97d15b 100644 --- a/src/primaite/game/agent/scripted_agents/__init__.py +++ b/src/primaite/game/agent/scripted_agents/__init__.py @@ -1 +1,6 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK + +from primaite.game.agent import interface +from primaite.game.agent.scripted_agents import abstract_tap, data_manipulation_bot, probabilistic_agent, random_agent + +__all__ = ("abstract_tap", "data_manipulation_bot", "interface", "probabilistic_agent", "random_agent") diff --git a/src/primaite/game/agent/scripted_agents/abstract_tap.py b/src/primaite/game/agent/scripted_agents/abstract_tap.py new file mode 100644 index 00000000..e6ddd546 --- /dev/null +++ b/src/primaite/game/agent/scripted_agents/abstract_tap.py @@ -0,0 +1,61 @@ +# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from __future__ import annotations + +import random +from abc import abstractmethod +from typing import Dict, List, Optional, Tuple + +from gymnasium.core import ObsType +from pydantic import Field + +from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent + +__all__ = "AbstractTAPAgent" + + +class AbstractTAPAgent(PeriodicAgent, identifier="AbstractTAP"): + """Base class for TAP agents to inherit from.""" + + config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema()) + next_execution_timestep: int = 0 + + class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema): + """Schema for the `agent_settings` part of the agent config.""" + + possible_starting_nodes: List[str] = Field(default_factory=list) + + class ConfigSchema(PeriodicAgent.ConfigSchema): + """Configuration schema for Abstract TAP agents.""" + + type: str = "AbstractTAP" + agent_settings: AbstractTAPAgent.AgentSettingsSchema = Field( + default_factory=lambda: AbstractTAPAgent.AgentSettingsSchema() + ) + + starting_node: Optional[str] = None + + @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) + + @abstractmethod + def setup_agent(self) -> None: + """Set up agent.""" + pass + + 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.config.agent_settings.variance, self.config.agent_settings.variance + ) + self.next_execution_timestep = timestep + random_timestep_increment + + 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 + self.starting_node = random.choice(self.config.agent_settings.possible_starting_nodes) + self.logger.debug(f"Selected starting node: {self.starting_node}") diff --git a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py index cf1b6ce8..a7558d42 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py @@ -1,31 +1,35 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -import random from typing import Dict, Tuple from gymnasium.core import ObsType +from pydantic import Field -from primaite.game.agent.interface import AbstractScriptedAgent +from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent + +__all__ = "DataManipulationAgent" -class DataManipulationAgent(AbstractScriptedAgent): +class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"): """Agent that uses a DataManipulationBot to perform an SQL injection attack.""" - next_execution_timestep: int = 0 - starting_node_idx: int = 0 + class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema): + """Schema for the `agent_settings` part of the agent config.""" - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.setup_agent() + target_application: str = "DataManipulationBot" - def _set_next_execution_timestep(self, timestep: int) -> None: - """Set the next execution timestep with a configured random variance. + class ConfigSchema(PeriodicAgent.ConfigSchema): + """Configuration Schema for DataManipulationAgent.""" - :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 + type: str = "RedDatabaseCorruptingAgent" + agent_settings: "DataManipulationAgent.AgentSettingsSchema" = Field( + default_factory=lambda: DataManipulationAgent.AgentSettingsSchema() ) - self.next_execution_timestep = timestep + random_timestep_increment + + config: "DataManipulationAgent.ConfigSchema" = Field(default_factory=lambda: DataManipulationAgent.ConfigSchema()) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._set_next_execution_timestep(timestep=self.config.agent_settings.start_step, variance=0) def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: """Waits until a specific timestep, then attempts to execute its data manipulation application. @@ -38,21 +42,14 @@ class DataManipulationAgent(AbstractScriptedAgent): :rtype: Tuple[str, Dict] """ if timestep < self.next_execution_timestep: - self.logger.debug(msg="Performing do NOTHING") - return "DONOTHING", {} + self.logger.debug(msg="Performing do nothing action") + return "do_nothing", {} - self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency) + self._set_next_execution_timestep( + timestep=timestep + self.config.agent_settings.frequency, variance=self.config.agent_settings.variance + ) self.logger.info(msg="Performing a data manipulation attack!") - 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) - self.logger.debug(msg=f"Select Start Node ID: {self.starting_node_idx}") + return "node_application_execute", { + "node_name": self.start_node, + "application_name": self.config.agent_settings.target_application, + } diff --git a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py index ce4d90d1..de643ed8 100644 --- a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py +++ b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py @@ -1,29 +1,28 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK """Agents with predefined behaviours.""" -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple import numpy as np import pydantic from gymnasium.core import ObsType +from numpy.random import Generator +from pydantic import Field -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 + +__all__ = "ProbabilisticAgent" -class ProbabilisticAgent(AbstractScriptedAgent): +class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): """Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" - class Settings(pydantic.BaseModel): - """Config schema for Probabilistic agent settings.""" + rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535))) - model_config = pydantic.ConfigDict(extra="forbid") - """Strict validation.""" - action_probabilities: Dict[int, float] + class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema): + """Schema for the `agent_settings` part of the agent config.""" + + action_probabilities: Dict[int, float] = None """Probability to perform each action in the action map. The sum of probabilities should sum to 1.""" - # 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 @@ -44,31 +43,20 @@ class ProbabilisticAgent(AbstractScriptedAgent): ) 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)}} + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """Configuration schema for Probabilistic Agent.""" - # The random number seed for np.random is dependent on whether a random number seed is set - # in the config file. If there is one it is processed by set_random_seed() in environment.py - # and as a consequence the the sequence of rng_seed's used here will be repeatable. - self.settings = ProbabilisticAgent.Settings(**settings) - rng_seed = np.random.randint(0, 65535) - self.rng = np.random.default_rng(rng_seed) + type: str = "ProbabilisticAgent" + agent_settings: "ProbabilisticAgent.AgentSettingsSchema" = Field( + default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema() + ) - # convert probabilities from - self.probabilities = np.asarray(list(self.settings.action_probabilities.values())) + config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema()) - super().__init__(agent_name, action_space, observation_space, reward_function) - self.logger.debug(f"ProbabilisticAgent RNG seed: {rng_seed}") + @property + def probabilities(self) -> Dict[str, int]: + """Convenience method to view the probabilities of the Agent.""" + return np.asarray(list(self.config.agent_settings.action_probabilities.values())) def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: """ diff --git a/src/primaite/game/agent/scripted_agents/random_agent.py b/src/primaite/game/agent/scripted_agents/random_agent.py index 59fe699a..721b5293 100644 --- a/src/primaite/game/agent/scripted_agents/random_agent.py +++ b/src/primaite/game/agent/scripted_agents/random_agent.py @@ -1,20 +1,27 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK import random -from typing import Dict, Optional, Tuple +from functools import cached_property +from typing import Dict, List, Tuple from gymnasium.core import ObsType -from pydantic import BaseModel +from pydantic import computed_field, Field, model_validator -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 + +__all__ = ("RandomAgent", "PeriodicAgent") -class RandomAgent(AbstractScriptedAgent): +class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"): """Agent that ignores its observation and acts completely at random.""" - def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: + config: "RandomAgent.ConfigSchema" = Field(default_factory=lambda: RandomAgent.ConfigSchema()) + + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """Configuration Schema for Random Agents.""" + + type: str = "RandomAgent" + + def get_action(self) -> Tuple[str, Dict]: """Sample the action space randomly. :param obs: Current observation for this agent, not used in RandomAgent @@ -27,41 +34,60 @@ class RandomAgent(AbstractScriptedAgent): return self.action_manager.get_action(self.action_manager.space.sample()) -class PeriodicAgent(AbstractScriptedAgent): +class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"): """Agent that does nothing most of the time, but executes application at regular intervals (with variance).""" - class Settings(BaseModel): - """Configuration values for when an agent starts performing actions.""" + config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema()) - start_step: int = 20 - "The timestep at which an agent begins performing it's actions." - start_variance: int = 5 - "Deviation around the start step." + class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema): + """Schema for the `agent_settings` part of the agent config.""" + + start_step: int = 5 + "The timestep at which an agent begins performing it's actions" frequency: int = 5 - "The number of timesteps to wait between performing actions." + "The number of timesteps to wait between performing actions" variance: int = 0 - "The amount the frequency can randomly change to." - max_executions: int = 999999 - "Maximum number of times the agent can execute its action." + "The amount the frequency can randomly change to" + possible_start_nodes: List[str] + target_application: str - def __init__( - self, - agent_name: str, - action_space: ActionManager, - observation_space: ObservationManager, - reward_function: RewardFunction, - settings: Optional[Settings] = None, - ) -> None: - """Initialise PeriodicAgent.""" - super().__init__( - agent_name=agent_name, - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, + @model_validator(mode="after") + def check_variance_lt_frequency(self) -> "PeriodicAgent.ConfigSchema": + """ + Make sure variance is equal to or lower than frequency. + + This is because the calculation for the next execution time is now + (frequency +- variance). + If variance were greater than frequency, sometimes the bracketed term would be negative + and the attack would never happen again. + """ + if self.variance >= self.frequency: + raise ValueError( + f"Agent start settings error: variance must be lower than frequency " + f"{self.variance=}, {self.frequency=}" + ) + return self + + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """Configuration Schema for Periodic Agent.""" + + type: str = "PeriodicAgent" + """Name of the agent.""" + agent_settings: "PeriodicAgent.AgentSettingsSchema" = Field( + default_factory=lambda: PeriodicAgent.AgentSettingsSchema() ) - self.settings = settings or PeriodicAgent.Settings() - self._set_next_execution_timestep(timestep=self.settings.start_step, variance=self.settings.start_variance) - self.num_executions = 0 + + max_executions: int = 999999 + "Maximum number of times the agent can execute its action." + num_executions: int = 0 + """Number of times the agent has executed an action.""" + next_execution_timestep: int = 0 + """Timestep of the next action execution by the agent.""" + + @computed_field + @cached_property + def start_node(self) -> str: + """On instantiation, randomly select a start node.""" + return random.choice(self.config.agent_settings.possible_start_nodes) def _set_next_execution_timestep(self, timestep: int, variance: int) -> None: """Set the next execution timestep with a configured random variance. @@ -76,9 +102,14 @@ class PeriodicAgent(AbstractScriptedAgent): def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: """Do nothing, unless the current timestep is the next execution timestep, in which case do the action.""" - if timestep == self.next_execution_timestep and self.num_executions < self.settings.max_executions: + if timestep == self.next_execution_timestep and self.num_executions < self.max_executions: self.num_executions += 1 - self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance) - return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0} + self._set_next_execution_timestep( + timestep + self.config.agent_settings.frequency, self.config.agent_settings.variance + ) + return "node_application_execute", { + "node_name": self.start_node, + "application_name": self.config.agent_settings.target_application, + } - return "DONOTHING", {} + return "do_nothing", {} diff --git a/src/primaite/game/agent/scripted_agents/tap001.py b/src/primaite/game/agent/scripted_agents/tap001.py deleted file mode 100644 index e9694a45..00000000 --- a/src/primaite/game/agent/scripted_agents/tap001.py +++ /dev/null @@ -1,78 +0,0 @@ -# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -import random -from typing import Dict, Tuple - -from gymnasium.core import ObsType - -from primaite.game.agent.interface import AbstractScriptedAgent - - -class TAP001(AbstractScriptedAgent): - """ - TAP001 | Mobile Malware -- Ransomware Variant. - - Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application) - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.setup_agent() - - next_execution_timestep: int = 0 - starting_node_idx: int = 0 - installed: bool = False - - 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 the ransomware application. - - This application acts a wrapper around the kill-chain, similar to green-analyst and - the previous UC2 data manipulation bot. - - :param obs: Current observation for this agent. - :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) - - if not self.installed: - self.installed = True - return "NODE_APPLICATION_INSTALL", { - "node_id": self.starting_node_idx, - "application_name": "RansomwareScript", - } - - 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) - for n, act in self.action_manager.action_map.items(): - if not act[0] == "NODE_APPLICATION_INSTALL": - continue - if act[1]["node_id"] == self.starting_node_idx: - self.ip_address = act[1]["ip_address"] - return - raise RuntimeError("TAP001 agent could not find database server ip address in action map") - - 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) diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index b9dc9c4d..009a376f 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -7,14 +7,8 @@ import numpy as np from pydantic import BaseModel, ConfigDict from primaite import DEFAULT_BANDWIDTH, getLogger -from primaite.game.agent.actions import ActionManager -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.agent.scripted_agents.random_agent import PeriodicAgent -from primaite.game.agent.scripted_agents.tap001 import TAP001 +from primaite.game.agent.interface import AbstractAgent, ProxyAgent +from primaite.game.agent.rewards import SharedReward from primaite.game.science import graph_has_cycle, topological_sort from primaite.simulator import SIM_OUTPUT from primaite.simulator.network.creation import NetworkNodeAdder @@ -44,7 +38,7 @@ from primaite.simulator.system.services.service import Service from primaite.simulator.system.services.terminal.terminal import Terminal from primaite.simulator.system.services.web_server.web_server import WebServer from primaite.simulator.system.software import Software -from primaite.utils.validation.ip_protocol import IPProtocol, PROTOCOL_LOOKUP +from primaite.utils.validation.ip_protocol import IPProtocol from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) @@ -258,6 +252,7 @@ class PrimaiteGame: net = sim.network simulation_config = cfg.get("simulation", {}) + defaults_config = cfg.get("defaults", {}) network_config = simulation_config.get("network", {}) airspace_cfg = network_config.get("airspace", {}) frequency_max_capacity_mbps_cfg = airspace_cfg.get("frequency_max_capacity_mbps", {}) @@ -283,6 +278,18 @@ class PrimaiteGame: _LOGGER.error(msg) raise ValueError(msg) + # TODO: handle simulation defaults more cleanly + if "node_start_up_duration" in defaults_config: + new_node.start_up_duration = defaults_config["node_startup_duration"] + if "node_shut_down_duration" in defaults_config: + new_node.shut_down_duration = defaults_config["node_shut_down_duration"] + if "node_scan_duration" in defaults_config: + new_node.node_scan_duration = defaults_config["node_scan_duration"] + if "folder_scan_duration" in defaults_config: + new_node.file_system._default_folder_scan_duration = defaults_config["folder_scan_duration"] + if "folder_restore_duration" in defaults_config: + new_node.file_system._default_folder_restore_duration = defaults_config["folder_restore_duration"] + if "users" in node_cfg and new_node.software_manager.software.get("UserManager"): user_manager: UserManager = new_node.software_manager.software["UserManager"] # noqa for user_cfg in node_cfg["users"]: @@ -315,12 +322,12 @@ class PrimaiteGame: if service_class is not None: _LOGGER.debug(f"installing {service_type} on node {new_node.config.hostname}") - new_node.software_manager.install(service_class, **service_cfg.get("options", {})) + new_node.software_manager.install(service_class) new_service = new_node.software_manager.software[service_class.__name__] # fixing duration for the service - if "fix_duration" in service_cfg.get("options", {}): - new_service.fixing_duration = service_cfg["options"]["fix_duration"] + if "fixing_duration" in service_cfg.get("options", {}): + new_service.config.fixing_duration = service_cfg["options"]["fixing_duration"] _set_software_listen_on_ports(new_service, service_cfg) # start the service @@ -329,6 +336,15 @@ class PrimaiteGame: msg = f"Configuration contains an invalid service type: {service_type}" _LOGGER.error(msg) raise ValueError(msg) + + # TODO: handle simulation defaults more cleanly + if "service_fix_duration" in defaults_config: + new_service.fixing_duration = defaults_config["service_fix_duration"] + if "service_restart_duration" in defaults_config: + new_service.restart_duration = defaults_config["service_restart_duration"] + if "service_install_duration" in defaults_config: + new_service.install_duration = defaults_config["service_install_duration"] + # service-dependent options if service_type == "DNSClient": if "options" in service_cfg: @@ -361,74 +377,20 @@ class PrimaiteGame: application_type = application_cfg["type"] if application_type in Application._registry: - new_node.software_manager.install(Application._registry[application_type]) + application_class = Application._registry[application_type] + application_options = application_cfg.get("options", {}) + application_options["type"] = application_type + new_node.software_manager.install(application_class, software_config=application_options) new_application = new_node.software_manager.software[application_type] # grab the instance - # fixing duration for the application - if "fix_duration" in application_cfg.get("options", {}): - new_application.fixing_duration = application_cfg["options"]["fix_duration"] else: msg = f"Configuration contains an invalid application type: {application_type}" _LOGGER.error(msg) raise ValueError(msg) - _set_software_listen_on_ports(new_application, application_cfg) - # run the application new_application.run() - if application_type == "DataManipulationBot": - if "options" in application_cfg: - opt = application_cfg["options"] - new_application.configure( - server_ip_address=IPv4Address(opt.get("server_ip")), - server_password=opt.get("server_password"), - 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")), - ) - elif application_type == "RansomwareScript": - if "options" in application_cfg: - opt = application_cfg["options"] - new_application.configure( - server_ip_address=IPv4Address(opt.get("server_ip")) if opt.get("server_ip") else None, - server_password=opt.get("server_password"), - payload=opt.get("payload", "ENCRYPT"), - ) - elif application_type == "DatabaseClient": - if "options" in application_cfg: - opt = application_cfg["options"] - new_application.configure( - server_ip_address=IPv4Address(opt.get("db_server_ip")), - server_password=opt.get("server_password"), - ) - elif application_type == "WebBrowser": - 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"] - new_application.configure( - target_ip_address=IPv4Address(opt.get("target_ip_address")), - target_port=PORT_LOOKUP[opt.get("target_port", "POSTGRES_SERVER")], - payload=opt.get("payload"), - repeat=bool(opt.get("repeat")), - port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")), - dos_intensity=float(opt.get("dos_intensity", "1.0")), - max_sessions=int(opt.get("max_sessions", "1000")), - ) - elif application_type == "C2Beacon": - if "options" in application_cfg: - opt = application_cfg["options"] - new_application.configure( - c2_server_ip_address=IPv4Address(opt.get("c2_server_ip_address")), - keep_alive_frequency=(opt.get("keep_alive_frequency", 5)), - masquerade_protocol=PROTOCOL_LOOKUP[ - (opt.get("masquerade_protocol", PROTOCOL_LOOKUP["TCP"])) - ], - masquerade_port=PORT_LOOKUP[(opt.get("masquerade_port", PORT_LOOKUP["HTTP"]))], - ) if "network_interfaces" in node_cfg: 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"])) @@ -470,76 +432,10 @@ class PrimaiteGame: agents_cfg = cfg.get("agents", []) for agent_cfg in agents_cfg: - agent_ref = agent_cfg["ref"] # noqa: F841 - agent_type = agent_cfg["type"] - action_space_cfg = agent_cfg["action_space"] - observation_space_cfg = agent_cfg["observation_space"] - reward_function_cfg = agent_cfg["reward_function"] - - # CREATE OBSERVATION SPACE - obs_space = ObservationManager.from_config(observation_space_cfg) - - # CREATE ACTION SPACE - action_space = ActionManager.from_config(game, action_space_cfg) - - # CREATE REWARD FUNCTION - reward_function = RewardFunction.from_config(reward_function_cfg) - - # CREATE AGENT - if agent_type == "ProbabilisticAgent": - # TODO: implement non-random agents and fix this parsing - 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, - settings=settings, - ) - elif agent_type == "PeriodicAgent": - settings = PeriodicAgent.Settings(**agent_cfg.get("settings", {})) - new_agent = PeriodicAgent( - agent_name=agent_cfg["ref"], - action_space=action_space, - observation_space=obs_space, - reward_function=reward_function, - settings=settings, - ) - - 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, - observation_space=obs_space, - reward_function=reward_function, - agent_settings=agent_settings, - ) - 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, - observation_space=obs_space, - reward_function=reward_function, - agent_settings=agent_settings, - ) - elif agent_type == "TAP001": - agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings")) - new_agent = TAP001( - agent_name=agent_cfg["ref"], - action_space=action_space, - observation_space=obs_space, - reward_function=reward_function, - agent_settings=agent_settings, - ) - else: - msg = f"Configuration error: {agent_type} is not a valid agent type." - _LOGGER.error(msg) - raise ValueError(msg) + new_agent = AbstractAgent.from_config(agent_cfg) game.agents[agent_cfg["ref"]] = new_agent + if isinstance(new_agent, ProxyAgent): + game.rl_agents[agent_cfg["ref"]] = new_agent # Validate that if any agents are sharing rewards, they aren't forming an infinite loop. game.setup_reward_sharing() diff --git a/src/primaite/notebooks/Action-masking.ipynb b/src/primaite/notebooks/Action-masking.ipynb index ba70f2b4..7fde0a49 100644 --- a/src/primaite/notebooks/Action-masking.ipynb +++ b/src/primaite/notebooks/Action-masking.ipynb @@ -11,6 +11,15 @@ "PrimAITE environments support action masking. The action mask shows which of the agent's actions are applicable with the current environment state. For example, a node can only be turned on if it is currently turned off." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -19,7 +28,7 @@ "source": [ "from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.config.load import data_manipulation_config_path\n", - "from prettytable import PrettyTable\n" + "from prettytable import PrettyTable" ] }, { @@ -195,7 +204,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -209,7 +218,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb index d2972fa9..66a684de 100644 --- a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb +++ b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb @@ -51,24 +51,15 @@ " - ref: CustomC2Agent\n", " team: RED\n", " type: ProxyAgent\n", - " observation_space: null\n", + "\n", " action_space:\n", - " action_list:\n", - " - type: DONOTHING\n", - " - type: NODE_APPLICATION_INSTALL\n", - " - type: NODE_APPLICATION_EXECUTE\n", - " - type: CONFIGURE_C2_BEACON\n", - " - type: C2_SERVER_RANSOMWARE_LAUNCH\n", - " - type: C2_SERVER_RANSOMWARE_CONFIGURE\n", - " - type: C2_SERVER_TERMINAL_COMMAND\n", - " - type: C2_SERVER_DATA_EXFILTRATE\n", " options:\n", " nodes:\n", " - node_name: web_server\n", - " applications: \n", + " applications:\n", " - application_name: C2Beacon\n", " - node_name: client_1\n", - " applications: \n", + " applications:\n", " - application_name: C2Server\n", " max_folders_per_node: 1\n", " max_files_per_folder: 1\n", @@ -82,15 +73,15 @@ " - 0.0.0.1\n", " action_map:\n", " 0:\n", - " action: DONOTHING\n", + " action: do_nothing\n", " options: {}\n", " 1:\n", - " action: NODE_APPLICATION_INSTALL\n", + " action: node_application_install\n", " options:\n", " node_id: 0\n", " application_name: C2Beacon\n", " 2:\n", - " action: CONFIGURE_C2_BEACON\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0\n", " config:\n", @@ -99,12 +90,12 @@ " masquerade_protocol:\n", " masquerade_port:\n", " 3:\n", - " action: NODE_APPLICATION_EXECUTE\n", + " action: node_application_execute\n", " options:\n", " node_id: 0\n", - " application_id: 0 \n", + " application_id: 0\n", " 4:\n", - " action: C2_SERVER_TERMINAL_COMMAND\n", + " action: c2_server_terminal_command\n", " options:\n", " node_id: 1\n", " ip_address:\n", @@ -112,20 +103,20 @@ " username: admin\n", " password: admin\n", " commands:\n", - " - \n", + " -\n", " - software_manager\n", " - application\n", " - install\n", " - RansomwareScript\n", " 5:\n", - " action: C2_SERVER_RANSOMWARE_CONFIGURE\n", + " action: c2_server_ransomware_configure\n", " options:\n", " node_id: 1\n", " config:\n", " server_ip_address: 192.168.1.14\n", " payload: ENCRYPT\n", " 6:\n", - " action: C2_SERVER_DATA_EXFILTRATE\n", + " action: c2_server_data_exfiltrate\n", " options:\n", " node_id: 1\n", " target_file_name: \"database.db\"\n", @@ -134,14 +125,14 @@ " target_ip_address: 192.168.1.14\n", " account:\n", " username: admin\n", - " password: admin \n", + " password: admin\n", "\n", " 7:\n", - " action: C2_SERVER_RANSOMWARE_LAUNCH\n", + " action: c2_server_ransomware_launch\n", " options:\n", " node_id: 1\n", " 8:\n", - " action: CONFIGURE_C2_BEACON\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0\n", " config:\n", @@ -150,7 +141,7 @@ " masquerade_protocol: TCP\n", " masquerade_port: DNS\n", " 9:\n", - " action: CONFIGURE_C2_BEACON\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0\n", " config:\n", @@ -177,7 +168,7 @@ " # removing all agents & adding the custom agent.\n", " cfg['agents'] = {}\n", " cfg['agents'] = c2_agent_yaml\n", - " \n", + "\n", "\n", "env = PrimaiteGymEnv(env_config=cfg)" ] @@ -222,7 +213,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Beacon Actions | NODE_APPLICATION_INSTALL\n", + "### **Command and Control** | C2 Beacon Actions | node_application_install\n", "\n", "The custom proxy red agent defined at the start of this notebook has been configured to install the C2 Beacon as action ``1`` in it's action map. \n", "\n", @@ -230,10 +221,6 @@ "\n", "```yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: NODE_APPLICATION_INSTALL\n", - " ...\n", " options:\n", " nodes: # Node List\n", " - node_name: web_server\n", @@ -243,7 +230,7 @@ " ...\n", " action_map:\n", " 1:\n", - " action: NODE_APPLICATION_INSTALL \n", + " action: node_application_install \n", " options:\n", " node_id: 0 # Index 0 at the node list.\n", " application_name: C2Beacon\n", @@ -265,7 +252,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Beacon Actions | CONFIGURE_C2_BEACON \n", + "### **Command and Control** | C2 Beacon Actions | configure_c2_beacon \n", "\n", "The custom proxy red agent defined at the start of this notebook can configure the C2 Beacon via action ``2`` in it's action map. \n", "\n", @@ -273,10 +260,6 @@ "\n", "```yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: CONFIGURE_C2_BEACON\n", - " ...\n", " options:\n", " nodes: # Node List\n", " - node_name: web_server\n", @@ -285,7 +268,7 @@ " action_map:\n", " ...\n", " 2:\n", - " action: CONFIGURE_C2_BEACON\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0 # Node Index\n", " config: # Further information about these config options can be found at the bottom of this notebook.\n", @@ -312,18 +295,14 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Beacon Actions | NODE_APPLICATION_EXECUTE\n", + "### **Command and Control** | C2 Beacon Actions | node_application_execute\n", "\n", - "The final action is ``NODE_APPLICATION_EXECUTE`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n", + "The final action is ``node_application_execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n", "\n", "The yaml snippet below shows all the relevant agent options for this action:\n", "\n", "```yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: NODE_APPLICATION_EXECUTE\n", - " ...\n", " options:\n", " nodes: # Node List\n", " - node_name: web_server\n", @@ -334,7 +313,7 @@ " action_map:\n", " ...\n", " 3:\n", - " action: NODE_APPLICATION_EXECUTE\n", + " action: node_application_execute\n", " options:\n", " node_id: 0\n", " application_id: 0\n", @@ -347,7 +326,7 @@ "metadata": {}, "outputs": [], "source": [ - "env.step(3) " + "env.step(3)" ] }, { @@ -390,10 +369,6 @@ "\n", "``` yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: C2_SERVER_TERMINAL_COMMAND\n", - " ...\n", " options:\n", " nodes: # Node List\n", " ...\n", @@ -441,7 +416,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | C2_SERVER_RANSOMWARE_CONFIGURE\n", + "### **Command and Control** | C2 Server Actions | c2_server_ransomware_configure\n", "\n", "Another action the C2 Server grants is the ability for a Red Agent to configure the RansomwareScript via the C2 Server rather than the note directly.\n", "\n", @@ -451,10 +426,6 @@ "\n", "``` yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: C2_SERVER_RANSOMWARE_CONFIGURE\n", - " ...\n", " options:\n", " nodes: # Node List\n", " ...\n", @@ -464,7 +435,7 @@ " ...\n", " action_map:\n", " 5:\n", - " action: C2_SERVER_RANSOMWARE_CONFIG\n", + " action: c2_server_ransomware_configure\n", " options:\n", " node_id: 1\n", " config:\n", @@ -497,9 +468,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | C2_SERVER_DATA_EXFILTRATE\n", + "### **Command and Control** | C2 Server Actions | c2_server_data_exfiltrate\n", "\n", - "The second to last action available is the ``C2_SERVER_DATA_EXFILTRATE`` which is indexed as action ``6`` in the action map.\n", + "The second to last action available is the ``c2_server_data_exfiltrate`` which is indexed as action ``6`` in the action map.\n", "\n", "This action can be used to exfiltrate a target file on a remote node to the C2 Beacon and the C2 Server's host file system via the ``FTP`` services.\n", "\n", @@ -507,10 +478,6 @@ "\n", "``` yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: C2_SERVER_DATA_EXFILTRATE\n", - " ...\n", " options:\n", " nodes: # Node List\n", " ...\n", @@ -520,7 +487,7 @@ " ...\n", " action_map:\n", " 6:\n", - " action: C2_SERVER_DATA_EXFILTRATE\n", + " action: c2_server_data_exfiltrate\n", " options:\n", " node_id: 1\n", " target_file_name: \"database.db\"\n", @@ -567,9 +534,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | C2_SERVER_RANSOMWARE_LAUNCH\n", + "### **Command and Control** | C2 Server Actions | c2_server_ransomware_launch\n", "\n", - "Finally, the last available action is for the C2_SERVER_RANSOMWARE_LAUNCH to start the ransomware script installed on the same node as the C2 beacon.\n", + "Finally, the last available action is for the c2_server_ransomware_launch to start the ransomware script installed on the same node as the C2 beacon.\n", "\n", "This action is indexed as action ``7``.\n", "\n", @@ -577,10 +544,6 @@ "\n", "``` yaml\n", " action_space:\n", - " action_list:\n", - " ...\n", - " - type: C2_SERVER_RANSOMWARE_LAUNCH\n", - " ...\n", " options:\n", " nodes: # Node List\n", " ...\n", @@ -590,7 +553,7 @@ " ...\n", " action_map:\n", " 7:\n", - " action: C2_SERVER_RANSOMWARE_LAUNCH\n", + " action: c2_server_ransomware_launch\n", " options:\n", " node_id: 1\n", "```\n" @@ -632,7 +595,7 @@ "metadata": {}, "outputs": [], "source": [ - "custom_blue_agent_yaml = \"\"\" \n", + "custom_blue_agent_yaml = \"\"\"\n", " - ref: defender\n", " team: BLUE\n", " type: ProxyAgent\n", @@ -715,28 +678,23 @@ " - type: \"NONE\"\n", " label: ICS\n", " options: {}\n", - " \n", + "\n", " action_space:\n", - " action_list:\n", - " - type: NODE_APPLICATION_REMOVE\n", - " - type: NODE_SHUTDOWN\n", - " - type: ROUTER_ACL_ADDRULE\n", - " - type: DONOTHING\n", " action_map:\n", " 0:\n", - " action: DONOTHING\n", + " action: do_nothing\n", " options: {}\n", " 1:\n", - " action: NODE_APPLICATION_REMOVE\n", + " action: node_application_remove\n", " options:\n", " node_id: 0\n", " application_name: C2Beacon\n", " 2:\n", - " action: NODE_SHUTDOWN\n", + " action: node_shutdown\n", " options:\n", " node_id: 0\n", " 3:\n", - " action: ROUTER_ACL_ADDRULE\n", + " action: router_acl_add_rule\n", " options:\n", " target_router: router_1\n", " position: 1\n", @@ -747,7 +705,7 @@ " dest_port_id: 2\n", " protocol_id: 1\n", " source_wildcard_id: 0\n", - " dest_wildcard_id: 0 \n", + " dest_wildcard_id: 0\n", "\n", "\n", " options:\n", @@ -796,7 +754,7 @@ " # removing all agents & adding the custom agent.\n", " cfg['agents'] = {}\n", " cfg['agents'] = custom_blue\n", - " \n", + "\n", "\n", "blue_env = PrimaiteGymEnv(env_config=cfg)" ] @@ -1121,7 +1079,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The code cell below uses the custom blue agent defined at the start of this section perform a NODE_APPLICATION_REMOVE on the C2 beacon:" + "The code cell below uses the custom blue agent defined at the start of this section perform a node_application_remove on the C2 beacon:" ] }, { @@ -1130,7 +1088,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Using CAOS ACTION: NODE_APPLICATION_REMOVE & capturing the OBS\n", + "# Using CAOS ACTION: node_application_remove & capturing the OBS\n", "post_blue_action_obs, _, _, _, _ = blue_env.step(1)" ] }, @@ -1216,7 +1174,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The code cell below uses the custom blue agent defined at the start of this section to perform a ``NODE_SHUT_DOWN`` action on the web server." + "The code cell below uses the custom blue agent defined at the start of this section to perform a ``node_shut_down`` action on the web server." ] }, { @@ -1225,7 +1183,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Using CAOS ACTION: NODE_SHUT_DOWN & capturing the OBS\n", + "# Using CAOS ACTION: node_shut_down & capturing the OBS\n", "post_blue_action_obs, _, _, _, _ = blue_env.step(2)" ] }, @@ -1306,7 +1264,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The code cell below uses the custom blue agent defined at the start of this section to perform a ROUTER_ACL_ADDRULE on router 1." + "The code cell below uses the custom blue agent defined at the start of this section to perform a router_acl_add_rule on router 1." ] }, { @@ -1315,7 +1273,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Using CAOS ACTION: ROUTER_ACL_ADDRULE & capturing the OBS\n", + "# Using CAOS ACTION: router_acl_add_rule & capturing the OBS\n", "post_blue_action_obs, _, _, _, _ = blue_env.step(3)" ] }, @@ -1429,11 +1387,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "As demonstrated earlier, red agents can use the ``CONFIGURE_C2_BEACON`` action to configure these settings mid episode through the configuration options:\n", + "As demonstrated earlier, red agents can use the ``configure_c2_beacon`` action to configure these settings mid episode through the configuration options:\n", "\n", "``` YAML\n", "...\n", - " action: CONFIGURE_C2_BEACON\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0\n", " config:\n", @@ -1468,7 +1426,7 @@ " # removing all agents & adding the custom agent.\n", " cfg['agents'] = {}\n", " cfg['agents'] = c2_agent_yaml\n", - " \n", + "\n", "\n", "c2_config_env = PrimaiteGymEnv(env_config=cfg)" ] @@ -1555,7 +1513,7 @@ "source": [ "for i in range(6):\n", " env.step(0)\n", - " \n", + "\n", "c2_server_1.show()" ] }, @@ -1676,7 +1634,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Comparing the OBS of the default frequency to a timestep frequency of 1 \n", + "# Comparing the OBS of the default frequency to a timestep frequency of 1\n", "for i in range(2):\n", " keep_alive_obs, _, _, _, _ = blue_config_env.step(0)\n", " display_obs_diffs(default_obs, keep_alive_obs, blue_config_env.game.step_counter)" @@ -1760,7 +1718,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Capturing default C2 Traffic \n", + "# Capturing default C2 Traffic\n", "for i in range(3):\n", " tcp_c2_obs, _, _, _, _ = blue_config_env.step(0)\n", "\n", diff --git a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb index dd5def9e..756fc44f 100644 --- a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb +++ b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb @@ -15,6 +15,15 @@ "*(For a full explanation of the Data Manipulation scenario, check out the data manipulation scenario notebook)*" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -67,9 +76,9 @@ " # parse the info dict form step output and write out what the red agent is doing\n", " red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n", " red_action = red_info.action\n", - " if red_action == 'DONOTHING':\n", + " if red_action == 'do_nothing':\n", " red_str = 'DO NOTHING'\n", - " elif red_action == 'NODE_APPLICATION_EXECUTE':\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" @@ -147,12 +156,7 @@ " 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", + " \n", " # The agent has access to the DataManipulationBoth on clients 1 and 2.\n", " options:\n", " nodes:\n", @@ -306,19 +310,9 @@ "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", + "# TODO:\n", "\"\"\")\n", + "#TODO 2869 fix\n", "\n", "with open(data_manipulation_config_path(), 'r') as f:\n", " cfg = yaml.safe_load(f)\n", @@ -444,7 +438,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": ".venv", "language": "python", "name": "python3" }, diff --git a/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb b/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb index 0460f771..dbc6f0c1 100644 --- a/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb +++ b/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb @@ -165,13 +165,13 @@ "\n", "| node_id | node name |\n", "|---------|------------------|\n", - "| 1 | domain_controller|\n", - "| 2 | web_server |\n", - "| 3 | database_server |\n", - "| 4 | backup_server |\n", - "| 5 | security_suite |\n", - "| 6 | client_1 |\n", - "| 7 | client_2 |\n", + "| 0 | domain_controller|\n", + "| 1 | web_server |\n", + "| 2 | database_server |\n", + "| 3 | backup_server |\n", + "| 4 | security_suite |\n", + "| 5 | client_1 |\n", + "| 6 | client_2 |\n", "\n", "Service 1 on node 2 (web_server) corresponds to the Web Server service. Other services are only there for padding to ensure that each node's observation space has the same shape. They are filled with zeroes.\n", "\n", @@ -371,6 +371,15 @@ "First, load the required modules" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -449,9 +458,9 @@ " # parse the info dict form step output and write out what the red agent is doing\n", " red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n", " red_action = red_info.action\n", - " if red_action == 'DONOTHING':\n", + " if red_action == 'do_nothing':\n", " red_str = 'DO NOTHING'\n", - " elif red_action == 'NODE_APPLICATION_EXECUTE':\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" @@ -547,7 +556,7 @@ "\n", "The reward will increase slightly as soon as the file finishes restoring. Then, the reward will increase to 0.9 when both green agents make successful requests.\n", "\n", - "Run the following cell until the green action is `NODE_APPLICATION_EXECUTE` for application 0, then the reward should increase. 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 increase. If you run it enough times, another red attack will happen and the reward will drop again." ] }, { diff --git a/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb b/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb index a832f3cc..f8691d7d 100644 --- a/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb +++ b/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb @@ -9,6 +9,15 @@ "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb b/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb similarity index 98% rename from src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb rename to src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb index c751edfd..fcda4dbd 100644 --- a/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb +++ b/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb @@ -201,7 +201,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -259,7 +259,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -396,7 +396,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] diff --git a/src/primaite/notebooks/Requests-and-Responses.ipynb b/src/primaite/notebooks/Requests-and-Responses.ipynb index da614c93..83aed07c 100644 --- a/src/primaite/notebooks/Requests-and-Responses.ipynb +++ b/src/primaite/notebooks/Requests-and-Responses.ipynb @@ -25,6 +25,15 @@ "Let's set up a minimal network simulation and send some requests to see how it works." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/src/primaite/notebooks/Terminal-Processing.ipynb b/src/primaite/notebooks/Terminal-Processing.ipynb index fdf405a7..9aa4e96a 100644 --- a/src/primaite/notebooks/Terminal-Processing.ipynb +++ b/src/primaite/notebooks/Terminal-Processing.ipynb @@ -18,6 +18,15 @@ "The Terminal service comes pre-installed on most Nodes (The exception being Switches, as these are currently dumb). " ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb b/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb index 19e95a95..76cab86a 100644 --- a/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb +++ b/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb @@ -18,6 +18,15 @@ "#### First, Import packages and read our config file." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -32,8 +41,6 @@ "from ray.rllib.algorithms.ppo import PPOConfig\n", "from primaite.session.ray_envs import PrimaiteRayMARLEnv\n", "\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/data_manipulation_marl.yaml', 'r') as f:\n", " cfg = yaml.safe_load(f)\n", "\n", diff --git a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb index dbe8871c..7252b046 100644 --- a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb @@ -11,6 +11,15 @@ "This notebook will demonstrate how to use PrimaiteRayEnv to train a basic PPO agent." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -95,7 +104,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -109,7 +118,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb index 892736fe..2b554475 100644 --- a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb @@ -18,6 +18,15 @@ "#### First, we import the inital packages and read in our configuration file." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -168,7 +177,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -182,7 +191,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.8" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Using-Episode-Schedules.ipynb b/src/primaite/notebooks/Using-Episode-Schedules.ipynb index cb06e0f9..d08ca67b 100644 --- a/src/primaite/notebooks/Using-Episode-Schedules.ipynb +++ b/src/primaite/notebooks/Using-Episode-Schedules.ipynb @@ -238,7 +238,7 @@ "### Episode 2\n", "When we reset the environment again, it moves onto episode 2, where it will bring in greens_1 and reds_1 for green and red agent definitions. Let's verify the agent names and that they take actions at the defined frequency.\n", "\n", - "Most green actions will be `NODE_APPLICATION_EXECUTE` while red will `DONOTHING` except at steps 10 and 20." + "Most green actions will be `node_application_execute` while red will `DONOTHING` except at steps 10 and 20." ] }, { @@ -269,7 +269,7 @@ "### Episode 3\n", "When we reset the environment again, it moves onto episode 3, where it will bring in greens_2 and reds_2 for green and red agent definitions. Let's verify the agent names and that they take actions at the defined frequency.\n", "\n", - "Now, green will perform `NODE_APPLICATION_EXECUTE` only 5% of the time, while red will perform `NODE_APPLICATION_EXECUTE` more frequently than before." + "Now, green will perform `node_application_execute` only 5% of the time, while red will perform `node_application_execute` more frequently than before." ] }, { diff --git a/src/primaite/notebooks/_package_data/uc2_attack.png b/src/primaite/notebooks/_package_data/uc2_attack.png index 8b8df5ce..03797d00 100644 Binary files a/src/primaite/notebooks/_package_data/uc2_attack.png and b/src/primaite/notebooks/_package_data/uc2_attack.png differ diff --git a/src/primaite/notebooks/_package_data/uc2_network.png b/src/primaite/notebooks/_package_data/uc2_network.png index 20fa43c9..10989201 100644 Binary files a/src/primaite/notebooks/_package_data/uc2_network.png and b/src/primaite/notebooks/_package_data/uc2_network.png differ diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb index 305cfd70..ad386f34 100644 --- a/src/primaite/notebooks/multi-processing.ipynb +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -18,6 +18,15 @@ "Import packages and read config file." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/src/primaite/session/environment.py b/src/primaite/session/environment.py index 8e608ede..b7a9a042 100644 --- a/src/primaite/session/environment.py +++ b/src/primaite/session/environment.py @@ -89,7 +89,7 @@ class PrimaiteGymEnv(gymnasium.Env): :return: Action mask :rtype: List[bool] """ - if not self.agent.action_masking: + if not self.agent.config.agent_settings.action_masking: return np.asarray([True] * len(self.agent.action_manager.action_map)) else: return self.game.action_mask(self._agent_name) diff --git a/src/primaite/session/ray_envs.py b/src/primaite/session/ray_envs.py index 33ba0540..16c85cb3 100644 --- a/src/primaite/session/ray_envs.py +++ b/src/primaite/session/ray_envs.py @@ -44,7 +44,7 @@ class PrimaiteRayMARLEnv(MultiAgentEnv): ) for agent_name in self._agent_ids: agent = self.game.rl_agents[agent_name] - if agent.action_masking: + if agent.config.agent_settings.action_masking: self.observation_space[agent_name] = spaces.Dict( { "action_mask": spaces.MultiBinary(agent.action_manager.space.n), @@ -143,7 +143,7 @@ class PrimaiteRayMARLEnv(MultiAgentEnv): unflat_space = agent.observation_manager.space unflat_obs = agent.observation_manager.current_observation obs = gymnasium.spaces.flatten(unflat_space, unflat_obs) - if agent.action_masking: + if agent.config.agent_settings.action_masking: all_obs[agent_name] = {"action_mask": self.game.action_mask(agent_name), "observations": obs} else: all_obs[agent_name] = obs @@ -168,7 +168,7 @@ class PrimaiteRayEnv(gymnasium.Env): self.env = PrimaiteGymEnv(env_config=env_config) # self.env.episode_counter -= 1 self.action_space = self.env.action_space - if self.env.agent.action_masking: + if self.env.agent.config.agent_settings.action_masking: self.observation_space = spaces.Dict( {"action_mask": spaces.MultiBinary(self.env.action_space.n), "observations": self.env.observation_space} ) @@ -178,7 +178,7 @@ class PrimaiteRayEnv(gymnasium.Env): def reset(self, *, seed: int = None, options: dict = None) -> Tuple[ObsType, Dict]: """Reset the environment.""" super().reset() # Ensure PRNG seed is set everywhere - if self.env.agent.action_masking: + if self.env.agent.config.agent_settings.action_masking: obs, *_ = self.env.reset(seed=seed) new_obs = {"action_mask": self.env.action_masks(), "observations": obs} return new_obs, *_ @@ -187,7 +187,7 @@ class PrimaiteRayEnv(gymnasium.Env): def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict]: """Perform a step in the environment.""" # if action masking is enabled, intercept the step method and add action mask to observation - if self.env.agent.action_masking: + if self.env.agent.config.agent_settings.action_masking: obs, *_ = self.env.step(action) new_obs = {"action_mask": self.game.action_mask(self.env._agent_name), "observations": obs} return new_obs, *_ diff --git a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb index 117ea019..7af8b98e 100644 --- a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb +++ b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb @@ -264,7 +264,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb index 4b620eee..58c07fe5 100644 --- a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb +++ b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb @@ -664,7 +664,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/simulator/file_system/file.py b/src/primaite/simulator/file_system/file.py index 57d01ec9..bad26a0a 100644 --- a/src/primaite/simulator/file_system/file.py +++ b/src/primaite/simulator/file_system/file.py @@ -130,8 +130,8 @@ class File(FileSystemItemABC): Return False if corruption is detected, otherwise True """ - warnings.warn("NODE_FILE_CHECKHASH is currently not implemented.") - self.sys_log.warning("NODE_FILE_CHECKHASH is currently not implemented.") + warnings.warn("node_file_checkhash is currently not implemented.") + self.sys_log.warning("node_file_checkhash is currently not implemented.") return False if self.deleted: diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 8ff4b6fb..54e649f2 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -30,6 +30,11 @@ class FileSystem(SimComponent): num_file_deletions: int = 0 "Number of file deletions in the current step." + _default_folder_scan_duration: Optional[int] = None + "Override default scan duration for folders" + _default_folder_restore_duration: Optional[int] = None + "Override default restore duration for folders" + def __init__(self, **kwargs): super().__init__(**kwargs) # Ensure a default root folder @@ -258,6 +263,11 @@ class FileSystem(SimComponent): name=folder.name, request_type=RequestType(func=folder._request_manager) ) self.folders[folder.uuid] = folder + # set the folder scan and restore durations. + if self._default_folder_scan_duration is not None: + folder.scan_duration = self._default_folder_scan_duration + if self._default_folder_restore_duration is not None: + folder.restore_duration = self._default_folder_restore_duration return folder def delete_folder(self, folder_name: str) -> bool: diff --git a/src/primaite/simulator/file_system/file_system_item_abc.py b/src/primaite/simulator/file_system/file_system_item_abc.py index 48b95d20..db51924c 100644 --- a/src/primaite/simulator/file_system/file_system_item_abc.py +++ b/src/primaite/simulator/file_system/file_system_item_abc.py @@ -43,6 +43,9 @@ def convert_size(size_bytes: int) -> str: class FileSystemItemHealthStatus(Enum): """Status of the FileSystemItem.""" + NONE = 0 + """File system item health status is not known.""" + GOOD = 1 """File/Folder is OK.""" @@ -72,7 +75,7 @@ class FileSystemItemABC(SimComponent): health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD "Actual status of the current FileSystemItem" - visible_health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD + visible_health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.NONE "Visible status of the current FileSystemItem" previous_hash: Optional[str] = None diff --git a/src/primaite/simulator/file_system/folder.py b/src/primaite/simulator/file_system/folder.py index ee0f3d01..3dd5d1ce 100644 --- a/src/primaite/simulator/file_system/folder.py +++ b/src/primaite/simulator/file_system/folder.py @@ -46,7 +46,7 @@ class Folder(FileSystemItemABC): :param sys_log: The SysLog instance to us to create system logs. """ super().__init__(**kwargs) - + self._scanned_this_step: bool = False self.sys_log.info(f"Created file /{self.name} (id: {self.uuid})") def _init_request_manager(self) -> RequestManager: @@ -83,6 +83,7 @@ class Folder(FileSystemItemABC): state = super().describe_state() state["files"] = {file.name: file.describe_state() for uuid, file in self.files.items()} state["deleted_files"] = {file.name: file.describe_state() for uuid, file in self.deleted_files.items()} + state["scanned_this_step"] = self._scanned_this_step return state def show(self, markdown: bool = False): @@ -135,7 +136,7 @@ class Folder(FileSystemItemABC): def pre_timestep(self, timestep: int) -> None: """Apply pre-timestep logic.""" super().pre_timestep(timestep) - + self._scanned_this_step = False for file in self.files.values(): file.pre_timestep(timestep) @@ -148,9 +149,17 @@ class Folder(FileSystemItemABC): for file_id in self.files: file = self.get_file_by_id(file_uuid=file_id) file.scan() - if file.visible_health_status == FileSystemItemHealthStatus.CORRUPT: - self.health_status = FileSystemItemHealthStatus.CORRUPT + # set folder health to worst file's health by generating a list of file healths. If no files, use 0 + self.health_status = FileSystemItemHealthStatus( + max( + [f.health_status.value for f in self.files.values()] + or [ + 0, + ] + ) + ) self.visible_health_status = self.health_status + self._scanned_this_step = True def _reveal_to_red_timestep(self) -> None: """Apply reveal to red timestep.""" @@ -387,8 +396,8 @@ class Folder(FileSystemItemABC): Return False if corruption is detected, otherwise True """ - warnings.warn("NODE_FOLDER_CHECKHASH is currently not implemented.") - self.sys_log.error("NODE_FOLDER_CHECKHASH is currently not implemented.") + warnings.warn("node_folder_checkhash is currently not implemented.") + self.sys_log.error("node_folder_checkhash is currently not implemented.") return False if self.deleted: diff --git a/src/primaite/simulator/network/creation.py b/src/primaite/simulator/network/creation.py index 94c45428..ebd17638 100644 --- a/src/primaite/simulator/network/creation.py +++ b/src/primaite/simulator/network/creation.py @@ -1,7 +1,7 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from abc import ABC, abstractmethod from ipaddress import IPv4Address -from typing import Any, ClassVar, Dict, Literal, Type +from typing import Any, ClassVar, Dict, Literal, Optional, Type from pydantic import BaseModel, model_validator @@ -49,7 +49,7 @@ class NetworkNodeAdder(BaseModel): _registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {} - def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None: """ Register a network node adder class. @@ -58,6 +58,8 @@ class NetworkNodeAdder(BaseModel): :raises ValueError: When attempting to register a name that is already reserved. """ super().__init_subclass__(**kwargs) + if identifier is None: + return if identifier in cls._registry: raise ValueError(f"Duplicate node adder {identifier}") cls._registry[identifier] = cls diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index dbe9705b..872b1bdf 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -824,7 +824,7 @@ class User(SimComponent): return self.model_dump() -class UserManager(Service): +class UserManager(Service, identifier="UserManager"): """ Manages users within the PrimAITE system, handling creation, authentication, and administration. @@ -833,11 +833,18 @@ class UserManager(Service): :param disabled_admins: A dictionary of currently disabled admin users by their usernames """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for UserManager.""" + + type: str = "UserManager" + + config: "UserManager.ConfigSchema" = Field(default_factory=lambda: UserManager.ConfigSchema()) + users: Dict[str, User] = {} def __init__(self, **kwargs): """ - Initializes a UserManager instanc. + Initializes a UserManager instance. :param username: The username for the default admin user :param password: The password for the default admin user @@ -1130,13 +1137,20 @@ class RemoteUserSession(UserSession): return state -class UserSessionManager(Service): +class UserSessionManager(Service, identifier="UserSessionManager"): """ Manages user sessions on a Node, including local and remote sessions. This class handles authentication, session management, and session timeouts for users interacting with the Node. """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for UserSessionManager.""" + + type: str = "UserSessionManager" + + config: "UserSessionManager.ConfigSchema" = Field(default_factory=lambda: UserSessionManager.ConfigSchema()) + local_session: Optional[UserSession] = None """The current local user session, if any.""" @@ -1554,7 +1568,6 @@ class Node(SimComponent, ABC): red_scan_countdown: int = 0 "Time steps until reveal to red scan is complete." - @classmethod def from_config(cls, config: Dict) -> "Node": """Create Node object from a given configuration dictionary.""" @@ -1564,7 +1577,7 @@ class Node(SimComponent, ABC): obj = cls(config=cls.ConfigSchema(**config)) return obj - def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: """ Register a node type. @@ -1572,10 +1585,10 @@ class Node(SimComponent, ABC): :type identifier: str :raises ValueError: When attempting to register an node with a name that is already allocated. """ - if identifier == "default": + super().__init_subclass__(**kwargs) + if identifier is None: return identifier = identifier.lower() - super().__init_subclass__(**kwargs) if identifier in cls._registry: raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.") cls._registry[identifier] = cls diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 1752c09a..05a47d7a 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -1,10 +1,12 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from __future__ import annotations -from abc import abstractmethod +from abc import ABC, abstractmethod from enum import Enum from typing import Any, ClassVar, Dict, Optional, Set, Type +from pydantic import Field + from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestPermissionValidator, RequestType from primaite.simulator.system.software import IOSoftware, SoftwareHealthState @@ -21,13 +23,20 @@ class ApplicationOperatingState(Enum): "The application is being installed or updated." -class Application(IOSoftware): +class Application(IOSoftware, ABC): """ Represents an Application in the simulation environment. Applications are user-facing programs that may perform input/output operations. """ + class ConfigSchema(IOSoftware.ConfigSchema, ABC): + """Config Schema for Application class.""" + + type: str + + config: ConfigSchema = Field(default_factory=lambda: Application.ConfigSchema()) + operating_state: ApplicationOperatingState = ApplicationOperatingState.CLOSED "The current operating state of the Application." execution_control_status: str = "manual" @@ -44,21 +53,36 @@ class Application(IOSoftware): _registry: ClassVar[Dict[str, Type["Application"]]] = {} """Registry of application types. Automatically populated when subclasses are defined.""" - def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: """ Register an application type. :param identifier: Uniquely specifies an application class by name. Used for finding items by config. - :type identifier: str + :type identifier: Optional[str] :raises ValueError: When attempting to register an application with a name that is already allocated. """ - if identifier == "default": - return super().__init_subclass__(**kwargs) + if identifier is None: + return if identifier in cls._registry: raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.") cls._registry[identifier] = cls + @classmethod + def from_config(cls, config: Dict) -> "Application": + """Create an application from a config dictionary. + + :param config: dict of options for application components constructor + :type config: dict + :return: The application component. + :rtype: Application + """ + if config["type"] not in cls._registry: + raise ValueError(f"Invalid Application type {config['type']}") + application_class = cls._registry[config["type"]] + application_object = application_class(config=application_class.ConfigSchema(**config)) + return application_object + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index f3ccc531..96130e16 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -6,7 +6,7 @@ from typing import Any, Dict, Optional, Union from uuid import uuid4 from prettytable import MARKDOWN, PrettyTable -from pydantic import BaseModel +from pydantic import BaseModel, Field from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -67,11 +67,19 @@ class DatabaseClient(Application, identifier="DatabaseClient"): Extends the Application class to provide functionality for connecting, querying, and disconnecting from a Database Service. It mainly operates over TCP protocol. - - :ivar server_ip_address: The IPv4 address of the Database Service server, defaults to None. """ + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for DatabaseClient.""" + + type: str = "DatabaseClient" + db_server_ip: Optional[IPV4Address] = None + server_password: Optional[str] = None + + config: ConfigSchema = Field(default_factory=lambda: DatabaseClient.ConfigSchema()) + server_ip_address: Optional[IPv4Address] = None + """The IPv4 address of the Database Service server, defaults to None.""" server_password: Optional[str] = None _query_success_tracker: Dict[str, bool] = {} """Keep track of connections that were established or verified during this step. Used for rewards.""" @@ -93,6 +101,8 @@ class DatabaseClient(Application, identifier="DatabaseClient"): kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"] kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"] super().__init__(**kwargs) + self.server_ip_address = self.config.db_server_ip + self.server_password = self.config.server_password def _init_request_manager(self) -> RequestManager: """ diff --git a/src/primaite/simulator/system/applications/nmap.py b/src/primaite/simulator/system/applications/nmap.py index f064eae3..3eeda4b6 100644 --- a/src/primaite/simulator/system/applications/nmap.py +++ b/src/primaite/simulator/system/applications/nmap.py @@ -3,7 +3,7 @@ from ipaddress import IPv4Address, IPv4Network from typing import Any, Dict, Final, List, Optional, Set, Tuple, Union from prettytable import PrettyTable -from pydantic import validate_call +from pydantic import Field, validate_call from primaite.interface.request import RequestResponse from primaite.simulator.core import RequestManager, RequestType, SimComponent @@ -52,6 +52,13 @@ class NMAP(Application, identifier="NMAP"): as ping scans to discover active hosts and port scans to detect open ports on those hosts. """ + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for NMAP.""" + + type: str = "NMAP" + + config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema()) + _active_port_scans: Dict[str, PortScanPayload] = {} _port_scan_responses: Dict[str, PortScanPayload] = {} diff --git a/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py b/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py index 4cd54d69..71a896bc 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py @@ -2,9 +2,9 @@ from abc import abstractmethod from enum import Enum from ipaddress import IPv4Address -from typing import Dict, Optional, Union +from typing import Dict, Optional, Set, Union -from pydantic import BaseModel, Field, validate_call +from pydantic import Field, validate_call from primaite.interface.request import RequestResponse from primaite.simulator.file_system.file_system import FileSystem, Folder @@ -45,10 +45,10 @@ class C2Payload(Enum): """C2 Input Command payload. Used by the C2 Server to send a command to the c2 beacon.""" OUTPUT = "output_command" - """C2 Output Command. Used by the C2 Beacon to send the results of a Input command to the c2 server.""" + """C2 Output Command. Used by the C2 Beacon to send the results of an Input command to the c2 server.""" -class AbstractC2(Application, identifier="AbstractC2"): +class AbstractC2(Application): """ An abstract command and control (c2) application. @@ -60,9 +60,25 @@ class AbstractC2(Application, identifier="AbstractC2"): Defaults to masquerading as HTTP (Port 80) via TCP. - Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite. + Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite. """ + class ConfigSchema(Application.ConfigSchema): + """Configuration for AbstractC2.""" + + keep_alive_frequency: int = Field(default=5, ge=1) + """The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon.""" + + masquerade_protocol: IPProtocol = Field(default=PROTOCOL_LOOKUP["TCP"]) + """The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP.""" + + masquerade_port: Port = Field(default=PORT_LOOKUP["HTTP"]) + """The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP.""" + + listen_on_ports: Set[Port] = {PORT_LOOKUP["HTTP"], PORT_LOOKUP["FTP"], PORT_LOOKUP["DNS"]} + + config: ConfigSchema = Field(default_factory=lambda: AbstractC2.ConfigSchema()) + c2_connection_active: bool = False """Indicates if the c2 server and c2 beacon are currently connected.""" @@ -75,19 +91,6 @@ class AbstractC2(Application, identifier="AbstractC2"): keep_alive_inactivity: int = 0 """Indicates how many timesteps since the last time the c2 application received a keep alive.""" - class _C2Opts(BaseModel): - """A Pydantic Schema for the different C2 configuration options.""" - - keep_alive_frequency: int = Field(default=5, ge=1) - """The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon.""" - - masquerade_protocol: IPProtocol = Field(default=PROTOCOL_LOOKUP["TCP"]) - """The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP.""" - - masquerade_port: Port = Field(default=PORT_LOOKUP["HTTP"]) - """The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP.""" - - c2_config: _C2Opts = _C2Opts() """ Holds the current configuration settings of the C2 Suite. @@ -100,6 +103,12 @@ class AbstractC2(Application, identifier="AbstractC2"): C2 beacon to reconfigure it's configuration settings. """ + def __init__(self, **kwargs): + """Initialise the C2 applications to by default listen for HTTP traffic.""" + kwargs["port"] = PORT_LOOKUP["NONE"] + kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"] + super().__init__(**kwargs) + def _craft_packet( self, c2_payload: C2Payload, c2_command: Optional[C2Command] = None, command_options: Optional[Dict] = {} ) -> C2Packet: @@ -118,13 +127,13 @@ class AbstractC2(Application, identifier="AbstractC2"): :type c2_command: C2Command. :param command_options: The relevant C2 Beacon parameters.F :type command_options: Dict - :return: Returns the construct C2Packet + :return: Returns the constructed C2Packet :rtype: C2Packet """ constructed_packet = C2Packet( - masquerade_protocol=self.c2_config.masquerade_protocol, - masquerade_port=self.c2_config.masquerade_port, - keep_alive_frequency=self.c2_config.keep_alive_frequency, + masquerade_protocol=self.config.masquerade_protocol, + masquerade_port=self.config.masquerade_port, + keep_alive_frequency=self.config.keep_alive_frequency, payload_type=c2_payload, command=c2_command, payload=command_options, @@ -140,13 +149,6 @@ class AbstractC2(Application, identifier="AbstractC2"): """ return super().describe_state() - def __init__(self, **kwargs): - """Initialise the C2 applications to by default listen for HTTP traffic.""" - kwargs["listen_on_ports"] = {PORT_LOOKUP["HTTP"], PORT_LOOKUP["FTP"], PORT_LOOKUP["DNS"]} - kwargs["port"] = PORT_LOOKUP["NONE"] - kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"] - super().__init__(**kwargs) - @property def _host_ftp_client(self) -> Optional[FTPClient]: """Return the FTPClient that is installed C2 Application's host. @@ -330,8 +332,8 @@ class AbstractC2(Application, identifier="AbstractC2"): if self.send( payload=keep_alive_packet, dest_ip_address=self.c2_remote_connection, - dest_port=self.c2_config.masquerade_port, - ip_protocol=self.c2_config.masquerade_protocol, + dest_port=self.config.masquerade_port, + ip_protocol=self.config.masquerade_protocol, session_id=session_id, ): # Setting the keep_alive_sent guard condition to True. This is used to prevent packet storms. @@ -340,8 +342,8 @@ class AbstractC2(Application, identifier="AbstractC2"): self.sys_log.info(f"{self.name}: Keep Alive sent to {self.c2_remote_connection}") self.sys_log.debug( f"{self.name}: Keep Alive sent to {self.c2_remote_connection} " - f"Masquerade Port: {self.c2_config.masquerade_port} " - f"Masquerade Protocol: {self.c2_config.masquerade_protocol} " + f"Masquerade Port: {self.config.masquerade_port} " + f"Masquerade Protocol: {self.config.masquerade_protocol} " ) return True else: @@ -376,15 +378,15 @@ class AbstractC2(Application, identifier="AbstractC2"): # Updating the C2 Configuration attribute. - self.c2_config.masquerade_port = payload.masquerade_port - self.c2_config.masquerade_protocol = payload.masquerade_protocol - self.c2_config.keep_alive_frequency = payload.keep_alive_frequency + self.config.masquerade_port = payload.masquerade_port + self.config.masquerade_protocol = payload.masquerade_protocol + self.config.keep_alive_frequency = payload.keep_alive_frequency self.sys_log.debug( f"{self.name}: C2 Config Resolved Config from Keep Alive:" - f"Masquerade Port: {self.c2_config.masquerade_port}" - f"Masquerade Protocol: {self.c2_config.masquerade_protocol}" - f"Keep Alive Frequency: {self.c2_config.keep_alive_frequency}" + f"Masquerade Port: {self.config.masquerade_port}" + f"Masquerade Protocol: {self.config.masquerade_protocol}" + f"Keep Alive Frequency: {self.config.keep_alive_frequency}" ) # This statement is intended to catch on the C2 Application that is listening for connection. @@ -410,8 +412,8 @@ class AbstractC2(Application, identifier="AbstractC2"): self.keep_alive_inactivity = 0 self.keep_alive_frequency = 5 self.c2_remote_connection = None - self.c2_config.masquerade_port = PORT_LOOKUP["HTTP"] - self.c2_config.masquerade_protocol = PROTOCOL_LOOKUP["TCP"] + self.config.masquerade_port = PORT_LOOKUP["HTTP"] + self.config.masquerade_protocol = PROTOCOL_LOOKUP["TCP"] @abstractmethod def _confirm_remote_connection(self, timestep: int) -> bool: diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py index b25eea6e..13918cd7 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py @@ -3,7 +3,7 @@ from ipaddress import IPv4Address from typing import Dict, Optional from prettytable import MARKDOWN, PrettyTable -from pydantic import validate_call +from pydantic import Field, validate_call from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -12,8 +12,9 @@ from primaite.simulator.system.applications.red_applications.c2 import ExfilOpts from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript from primaite.simulator.system.services.terminal.terminal import Terminal, TerminalClientConnection -from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP -from primaite.utils.validation.port import PORT_LOOKUP +from primaite.utils.validation.ip_protocol import IPProtocol, PROTOCOL_LOOKUP +from primaite.utils.validation.ipv4_address import IPV4Address +from primaite.utils.validation.port import Port, PORT_LOOKUP class C2Beacon(AbstractC2, identifier="C2Beacon"): @@ -32,15 +33,30 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): 2. Leveraging the terminal application to execute requests (dependent on the command given) 3. Sending the RequestResponse back to the C2 Server (Command output) - Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite. + Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite. """ + class ConfigSchema(AbstractC2.ConfigSchema): + """ConfigSchema for C2Beacon.""" + + type: str = "C2Beacon" + c2_server_ip_address: Optional[IPV4Address] = None + keep_alive_frequency: int = 5 + masquerade_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"] + masquerade_port: Port = PORT_LOOKUP["HTTP"] + + config: ConfigSchema = Field(default_factory=lambda: C2Beacon.ConfigSchema()) + keep_alive_attempted: bool = False """Indicates if a keep alive has been attempted to be sent this timestep. Used to prevent packet storms.""" terminal_session: TerminalClientConnection = None "The currently in use terminal session." + def __init__(self, **kwargs): + kwargs["name"] = "C2Beacon" + super().__init__(**kwargs) + @property def _host_terminal(self) -> Optional[Terminal]: """Return the Terminal that is installed on the same machine as the C2 Beacon.""" @@ -119,10 +135,6 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): rm.add_request("configure", request_type=RequestType(func=_configure)) return rm - def __init__(self, **kwargs): - kwargs["name"] = "C2Beacon" - super().__init__(**kwargs) - # Configure is practically setter method for the ``c2.config`` attribute that also ties into the request manager. @validate_call def configure( @@ -146,7 +158,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): masquerade_port | What port should the C2 traffic use? (TCP or UDP) These configuration options are used to reassign the fields in the inherited inner class - ``c2_config``. + ``config``. If a connection is already in progress then this method also sends a keep alive to the C2 Server in order for the C2 Server to sync with the new configuration settings. @@ -162,9 +174,9 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): :return: Returns True if the configuration was successful, False otherwise. """ self.c2_remote_connection = IPv4Address(c2_server_ip_address) - self.c2_config.keep_alive_frequency = keep_alive_frequency - self.c2_config.masquerade_port = masquerade_port - self.c2_config.masquerade_protocol = masquerade_protocol + self.config.keep_alive_frequency = keep_alive_frequency + self.config.masquerade_port = masquerade_port + self.config.masquerade_protocol = masquerade_protocol self.sys_log.info( f"{self.name}: Configured {self.name} with remote C2 server connection: {c2_server_ip_address=}." ) @@ -263,14 +275,12 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): if self.send( payload=output_packet, dest_ip_address=self.c2_remote_connection, - dest_port=self.c2_config.masquerade_port, - ip_protocol=self.c2_config.masquerade_protocol, + dest_port=self.config.masquerade_port, + ip_protocol=self.config.masquerade_protocol, session_id=session_id, ): self.sys_log.info(f"{self.name}: Command output sent to {self.c2_remote_connection}") - self.sys_log.debug( - f"{self.name}: on {self.c2_config.masquerade_port} via {self.c2_config.masquerade_protocol}" - ) + self.sys_log.debug(f"{self.name}: on {self.config.masquerade_port} via {self.config.masquerade_protocol}") return True else: self.sys_log.warning( @@ -562,7 +572,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): :rtype bool: """ self.keep_alive_attempted = False # Resetting keep alive sent. - if self.keep_alive_inactivity == self.c2_config.keep_alive_frequency: + if self.keep_alive_inactivity == self.config.keep_alive_frequency: self.sys_log.info( f"{self.name}: Attempting to Send Keep Alive to {self.c2_remote_connection} at timestep {timestep}." ) @@ -627,9 +637,9 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): self.c2_connection_active, self.c2_remote_connection, self.keep_alive_inactivity, - self.c2_config.keep_alive_frequency, - self.c2_config.masquerade_protocol, - self.c2_config.masquerade_port, + self.config.keep_alive_frequency, + self.config.masquerade_protocol, + self.config.masquerade_port, ] ) print(table) diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py index 654b86e7..9d2097e9 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py @@ -2,7 +2,7 @@ from typing import Dict, Optional from prettytable import MARKDOWN, PrettyTable -from pydantic import validate_call +from pydantic import Field, validate_call from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -31,9 +31,16 @@ class C2Server(AbstractC2, identifier="C2Server"): 1. Sending commands to the C2 Beacon. (Command input) 2. Parsing terminal RequestResponses back to the Agent. - Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite. + Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite. """ + class ConfigSchema(AbstractC2.ConfigSchema): + """ConfigSchema for C2Server.""" + + type: str = "C2Server" + + config: ConfigSchema = Field(default_factory=lambda: C2Server.ConfigSchema()) + current_command_output: RequestResponse = None """The Request Response by the last command send. This attribute is updated by the method _handle_command_output.""" @@ -251,8 +258,8 @@ class C2Server(AbstractC2, identifier="C2Server"): payload=command_packet, dest_ip_address=self.c2_remote_connection, session_id=self.c2_session.uuid, - dest_port=self.c2_config.masquerade_port, - ip_protocol=self.c2_config.masquerade_protocol, + dest_port=self.config.masquerade_port, + ip_protocol=self.config.masquerade_protocol, ): self.sys_log.info(f"{self.name}: Successfully sent {given_command}.") self.sys_log.info(f"{self.name}: Awaiting command response {given_command}.") @@ -334,11 +341,11 @@ class C2Server(AbstractC2, identifier="C2Server"): :return: Returns False if the C2 beacon is considered dead. Otherwise True. :rtype bool: """ - if self.keep_alive_inactivity > self.c2_config.keep_alive_frequency: + if self.keep_alive_inactivity > self.config.keep_alive_frequency: self.sys_log.info(f"{self.name}: C2 Beacon connection considered dead due to inactivity.") self.sys_log.debug( f"{self.name}: Did not receive expected keep alive connection from {self.c2_remote_connection}" - f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.c2_config.keep_alive_frequency}" + f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.config.keep_alive_frequency}" f"{self.name}: Last Keep Alive received at {(timestep - self.keep_alive_inactivity)}" ) self._reset_c2_connection() @@ -389,8 +396,8 @@ class C2Server(AbstractC2, identifier="C2Server"): [ self.c2_connection_active, self.c2_remote_connection, - self.c2_config.masquerade_protocol, - self.c2_config.masquerade_port, + self.config.masquerade_protocol, + self.config.masquerade_port, ] ) print(table) diff --git a/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py b/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py index 0423087e..392cdfba 100644 --- a/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py @@ -3,6 +3,8 @@ from enum import IntEnum from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.game.science import simulate_trial from primaite.interface.request import RequestResponse @@ -10,6 +12,7 @@ from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP +from primaite.utils.validation.ipv4_address import IPV4Address from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) @@ -40,6 +43,18 @@ class DataManipulationAttackStage(IntEnum): class DataManipulationBot(Application, identifier="DataManipulationBot"): """A bot that simulates a script which performs a SQL injection attack.""" + class ConfigSchema(Application.ConfigSchema): + """Configuration schema for DataManipulationBot.""" + + type: str = "DataManipulationBot" + server_ip: Optional[IPV4Address] = None + server_password: Optional[str] = None + payload: str = "DELETE" + port_scan_p_of_success: float = 0.1 + data_manipulation_p_of_success: float = 0.1 + + config: "DataManipulationBot.ConfigSchema" = Field(default_factory=lambda: DataManipulationBot.ConfigSchema()) + payload: Optional[str] = None port_scan_p_of_success: float = 0.1 data_manipulation_p_of_success: float = 0.1 @@ -56,6 +71,12 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"): super().__init__(**kwargs) self._db_connection: Optional[DatabaseClientConnection] = None + self.server_ip_address = self.config.server_ip + self.server_password = self.config.server_password + self.payload = self.config.payload + self.port_scan_p_of_success = self.config.port_scan_p_of_success + self.data_manipulation_p_of_success = self.config.data_manipulation_p_of_success + def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. diff --git a/src/primaite/simulator/system/applications/red_applications/dos_bot.py b/src/primaite/simulator/system/applications/red_applications/dos_bot.py index 99c4acb3..ea7a4d8d 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -3,11 +3,14 @@ from enum import IntEnum from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.game.science import simulate_trial from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.system.applications.database_client import DatabaseClient +from primaite.utils.validation.ipv4_address import IPV4Address from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) @@ -32,6 +35,20 @@ class DoSAttackStage(IntEnum): class DoSBot(DatabaseClient, identifier="DoSBot"): """A bot that simulates a Denial of Service attack.""" + class ConfigSchema(DatabaseClient.ConfigSchema): + """ConfigSchema for DoSBot.""" + + type: str = "DoSBot" + target_ip_address: Optional[IPV4Address] = None + target_port: Port = PORT_LOOKUP["POSTGRES_SERVER"] + payload: Optional[str] = None + repeat: bool = False + port_scan_p_of_success: float = 0.1 + dos_intensity: float = 1.0 + max_sessions: int = 1000 + + config: "DoSBot.ConfigSchema" = Field(default_factory=lambda: DoSBot.ConfigSchema()) + target_ip_address: Optional[IPv4Address] = None """IP address of the target service.""" @@ -56,7 +73,13 @@ class DoSBot(DatabaseClient, identifier="DoSBot"): def __init__(self, **kwargs): super().__init__(**kwargs) self.name = "DoSBot" - self.max_sessions = 1000 # override normal max sessions + self.target_ip_address = self.config.target_ip_address + self.target_port = self.config.target_port + self.payload = self.config.payload + self.repeat = self.config.repeat + self.port_scan_p_of_success = self.config.port_scan_p_of_success + self.dos_intensity = self.config.dos_intensity + self.max_sessions = self.config.max_sessions def _init_request_manager(self) -> RequestManager: """ diff --git a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py index 3a8ac5ae..114d5716 100644 --- a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py +++ b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py @@ -3,12 +3,14 @@ from ipaddress import IPv4Address from typing import Dict, Optional from prettytable import MARKDOWN, PrettyTable +from pydantic import Field from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP +from primaite.utils.validation.ipv4_address import IPV4Address from primaite.utils.validation.port import PORT_LOOKUP @@ -18,6 +20,16 @@ class RansomwareScript(Application, identifier="RansomwareScript"): :ivar payload: The attack stage query payload. (Default ENCRYPT) """ + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for RansomwareScript.""" + + type: str = "RansomwareScript" + server_ip: Optional[IPV4Address] = None + server_password: Optional[str] = None + payload: str = "ENCRYPT" + + config: "RansomwareScript.ConfigSchema" = Field(default_factory=lambda: RansomwareScript.ConfigSchema()) + server_ip_address: Optional[IPv4Address] = None """IP address of node which hosts the database.""" server_password: Optional[str] = None @@ -32,6 +44,9 @@ class RansomwareScript(Application, identifier="RansomwareScript"): super().__init__(**kwargs) self._db_connection: Optional[DatabaseClientConnection] = None + self.server_ip_address = self.config.server_ip + self.server_password = self.config.server_password + self.payload = self.config.payload def describe_state(self) -> Dict: """ diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index ff185e2a..49f303b5 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address from typing import Dict, List, Optional from urllib.parse import urlparse -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from primaite import getLogger from primaite.interface.request import RequestResponse @@ -30,7 +30,13 @@ class WebBrowser(Application, identifier="WebBrowser"): The application requests and loads web pages using its domain name and requesting IP addresses using DNS. """ - target_url: Optional[str] = None + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for WebBrowser.""" + + type: str = "WebBrowser" + target_url: Optional[str] = None + + config: "WebBrowser.ConfigSchema" = Field(default_factory=lambda: WebBrowser.ConfigSchema()) domain_name_ip_address: Optional[IPv4Address] = None "The IP address of the domain name for the webpage." @@ -86,7 +92,7 @@ class WebBrowser(Application, identifier="WebBrowser"): :param: url: The address of the web page the browser requests :type: url: str """ - url = url or self.target_url + url = url or self.config.target_url if not self._can_perform_action(): return False diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index f0ee6f7c..ddb30a3b 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -106,7 +106,7 @@ class SoftwareManager: return True return False - def install(self, software_class: Type[IOSoftware], **install_kwargs): + def install(self, software_class: Type[IOSoftware], software_config: Optional[IOSoftware.ConfigSchema] = None): """ Install an Application or Service. @@ -115,13 +115,22 @@ class SoftwareManager: if software_class in self._software_class_to_name_map: self.sys_log.warning(f"Cannot install {software_class} as it is already installed") return - software = software_class( - software_manager=self, - sys_log=self.sys_log, - file_system=self.file_system, - dns_server=self.dns_server, - **install_kwargs, - ) + if software_config is None: + software = software_class( + software_manager=self, + sys_log=self.sys_log, + file_system=self.file_system, + dns_server=self.dns_server, + ) + else: + software = software_class( + software_manager=self, + sys_log=self.sys_log, + file_system=self.file_system, + dns_server=self.dns_server, + config=software_config, + ) + software.parent = self.node if isinstance(software, Application): self.node.applications[software.uuid] = software diff --git a/src/primaite/simulator/system/services/arp/arp.py b/src/primaite/simulator/system/services/arp/arp.py index 31938e83..bbeec301 100644 --- a/src/primaite/simulator/system/services/arp/arp.py +++ b/src/primaite/simulator/system/services/arp/arp.py @@ -5,6 +5,7 @@ from abc import abstractmethod from typing import Any, Dict, Optional, Union from prettytable import MARKDOWN, PrettyTable +from pydantic import Field from primaite.simulator.network.hardware.base import NetworkInterface from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket @@ -14,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address from primaite.utils.validation.port import PORT_LOOKUP -class ARP(Service): +class ARP(Service, identifier="ARP"): """ The ARP (Address Resolution Protocol) Service. @@ -22,6 +23,13 @@ class ARP(Service): sends ARP requests and replies, and processes incoming ARP packets. """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for ARP.""" + + type: str = "ARP" + + config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema()) + arp: Dict[IPV4Address, ARPEntry] = {} def __init__(self, **kwargs): diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 3a5f5b31..4ba4c4d4 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -3,6 +3,8 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Literal, Optional, Union from uuid import uuid4 +from pydantic import Field + from primaite import getLogger from primaite.simulator.file_system.file_system import File from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus @@ -17,13 +19,21 @@ from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) -class DatabaseService(Service): +class DatabaseService(Service, identifier="DatabaseService"): """ A class for simulating a generic SQL Server service. This class inherits from the `Service` class and provides methods to simulate a SQL database. """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for DatabaseService.""" + + type: str = "DatabaseService" + backup_server_ip: Optional[IPv4Address] = None + + config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema()) + password: Optional[str] = None """Password that needs to be provided by clients if they want to connect to the DatabaseService.""" @@ -42,6 +52,7 @@ class DatabaseService(Service): kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"] super().__init__(**kwargs) self._create_db_file() + self.backup_server_ip = self.config.backup_server_ip def install(self): """ diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 02cf54ae..0756eb05 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/src/primaite/simulator/system/services/dns/dns_client.py @@ -2,6 +2,8 @@ from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest from primaite.simulator.system.core.software_manager import SoftwareManager @@ -12,9 +14,15 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) -class DNSClient(Service): +class DNSClient(Service, identifier="DNSClient"): """Represents a DNS Client as a Service.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for DNSClient.""" + + type: str = "DNSClient" + + config: "DNSClient.ConfigSchema" = Field(default_factory=lambda: DNSClient.ConfigSchema()) dns_cache: Dict[str, IPv4Address] = {} "A dict of known mappings between domain/URLs names and IPv4 addresses." dns_server: Optional[IPv4Address] = None diff --git a/src/primaite/simulator/system/services/dns/dns_server.py b/src/primaite/simulator/system/services/dns/dns_server.py index b7c9a42c..3a1c0e18 100644 --- a/src/primaite/simulator/system/services/dns/dns_server.py +++ b/src/primaite/simulator/system/services/dns/dns_server.py @@ -3,6 +3,7 @@ from ipaddress import IPv4Address from typing import Any, Dict, Optional from prettytable import MARKDOWN, PrettyTable +from pydantic import Field from primaite import getLogger from primaite.simulator.network.protocols.dns import DNSPacket @@ -13,9 +14,17 @@ from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) -class DNSServer(Service): +class DNSServer(Service, identifier="DNSServer"): """Represents a DNS Server as a Service.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for DNSServer.""" + + type: str = "DNSServer" + domain_mapping: dict = {} + + config: "DNSServer.ConfigSchema" = Field(default_factory=lambda: DNSServer.ConfigSchema()) + dns_table: Dict[str, IPv4Address] = {} "A dict of mappings between domain names and IPv4 addresses." diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 9c7b91ce..82875b97 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -2,6 +2,8 @@ from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -9,20 +11,28 @@ from primaite.simulator.file_system.file_system import File from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC +from primaite.simulator.system.services.service import Service from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) -class FTPClient(FTPServiceABC): +class FTPClient(FTPServiceABC, identifier="FTPClient"): """ A class for simulating an FTP client service. - This class inherits from the `Service` class and provides methods to emulate FTP + This class inherits from the `FTPServiceABC` class and provides methods to emulate FTP RFC 959: https://datatracker.ietf.org/doc/html/rfc959 """ + config: "FTPClient.ConfigSchema" = Field(default_factory=lambda: FTPClient.ConfigSchema()) + + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for FTPClient.""" + + type: str = "FTPClient" + def __init__(self, **kwargs): kwargs["name"] = "FTPClient" kwargs["port"] = PORT_LOOKUP["FTP"] @@ -108,6 +118,7 @@ class FTPClient(FTPServiceABC): session_id: Optional[str] = None, is_reattempt: Optional[bool] = False, ) -> bool: + self._active = True """ Connects the client to a given FTP server. @@ -164,6 +175,7 @@ class FTPClient(FTPServiceABC): :param: is_reattempt: Set to True if attempt to disconnect from FTP Server has been attempted. Default False. :type: is_reattempt: Optional[bool] """ + self._active = True # send a disconnect request payload to FTP server payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.QUIT) software_manager: SoftwareManager = self.software_manager @@ -209,6 +221,7 @@ class FTPClient(FTPServiceABC): :param: session_id: The id of the session :type: session_id: Optional[str] """ + self._active = True # check if the file to transfer exists on the client file_to_transfer: File = self.file_system.get_file(folder_name=src_folder_name, file_name=src_file_name) if not file_to_transfer: @@ -266,6 +279,7 @@ class FTPClient(FTPServiceABC): :param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"]. :type: dest_port: Optional[int] """ + self._active = True # check if FTP is currently connected to IP self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port) @@ -317,6 +331,7 @@ class FTPClient(FTPServiceABC): This helps prevent an FTP request loop - FTP client and servers can exist on the same node. """ + self._active = True if not self._can_perform_action(): return False diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index 9ce7d658..054bfe15 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -1,26 +1,36 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from typing import Any, Optional +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC +from primaite.simulator.system.services.service import Service from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP _LOGGER = getLogger(__name__) -class FTPServer(FTPServiceABC): +class FTPServer(FTPServiceABC, identifier="FTPServer"): """ A class for simulating an FTP server service. - This class inherits from the `Service` class and provides methods to emulate FTP + This class inherits from the `FTPServiceABC` class and provides methods to emulate FTP RFC 959: https://datatracker.ietf.org/doc/html/rfc959 """ + config: "FTPServer.ConfigSchema" = Field(default_factory=lambda: FTPServer.ConfigSchema()) + server_password: Optional[str] = None """Password needed to connect to FTP server. Default is None.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for FTPServer.""" + + type: str = "FTPServer" + def __init__(self, **kwargs): kwargs["name"] = "FTPServer" kwargs["port"] = PORT_LOOKUP["FTP"] diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index 52f451e1..13acda70 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -3,9 +3,11 @@ from abc import ABC from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import StrictBool + from primaite.simulator.file_system.file_system import File from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode -from primaite.simulator.system.services.service import Service +from primaite.simulator.system.services.service import Service, ServiceOperatingState from primaite.utils.validation.port import Port @@ -16,9 +18,22 @@ class FTPServiceABC(Service, ABC): Contains shared methods between both classes. """ + _active: StrictBool = False + """Flag that is True on timesteps where service transmits data and False when idle. Used for describe_state.""" + + def pre_timestep(self, timestep: int) -> None: + """When a new timestep begins, clear the _active attribute.""" + self._active = False + return super().pre_timestep(timestep) + def describe_state(self) -> Dict: """Returns a Dict of the FTPService state.""" - return super().describe_state() + state = super().describe_state() + + # override so that the service is shows as running only if actively transmitting data this timestep + if self.operating_state == ServiceOperatingState.RUNNING and not self._active: + state["operating_state"] = ServiceOperatingState.STOPPED.value + return state def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket: """ @@ -29,6 +44,7 @@ class FTPServiceABC(Service, ABC): :param: session_id: session ID linked to the FTP Packet. Optional. :type: session_id: Optional[str] """ + self._active = True if payload.ftp_command is not None: self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.") @@ -51,6 +67,7 @@ class FTPServiceABC(Service, ABC): :param: payload: The FTP Packet that contains the file data :type: FTPPacket """ + self._active = True try: file_name = payload.ftp_command_args["dest_file_name"] folder_name = payload.ftp_command_args["dest_folder_name"] @@ -106,6 +123,7 @@ class FTPServiceABC(Service, ABC): :param: is_response: is true if the data being sent is in response to a request. Default False. :type: is_response: bool """ + self._active = True # send STOR request payload: FTPPacket = FTPPacket( ftp_command=FTPCommand.STOR, @@ -135,6 +153,7 @@ class FTPServiceABC(Service, ABC): :param: payload: The FTP Packet that contains the file data :type: FTPPacket """ + self._active = True try: # find the file file_name = payload.ftp_command_args["src_file_name"] @@ -181,6 +200,7 @@ class FTPServiceABC(Service, ABC): :return: True if successful, False otherwise. """ + self._active = True self.sys_log.info(f"{self.name}: Sending FTP {payload.ftp_command.name} {payload.ftp_command_args}") return super().send( diff --git a/src/primaite/simulator/system/services/icmp/icmp.py b/src/primaite/simulator/system/services/icmp/icmp.py index 933d0591..7f626945 100644 --- a/src/primaite/simulator/system/services/icmp/icmp.py +++ b/src/primaite/simulator/system/services/icmp/icmp.py @@ -3,6 +3,8 @@ import secrets from ipaddress import IPv4Address from typing import Any, Dict, Optional, Tuple, Union +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.hardware.base import NetworkInterface from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType @@ -14,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) -class ICMP(Service): +class ICMP(Service, identifier="ICMP"): """ The Internet Control Message Protocol (ICMP) service. @@ -22,6 +24,13 @@ class ICMP(Service): network diagnostics, notably the ping command. """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for ICMP.""" + + type: str = "ICMP" + + config: "ICMP.ConfigSchema" = Field(default_factory=lambda: ICMP.ConfigSchema()) + request_replies: Dict = {} def __init__(self, **kwargs): diff --git a/src/primaite/simulator/system/services/ntp/ntp_client.py b/src/primaite/simulator/system/services/ntp/ntp_client.py index 9606c61f..fb470faf 100644 --- a/src/primaite/simulator/system/services/ntp/ntp_client.py +++ b/src/primaite/simulator/system/services/ntp/ntp_client.py @@ -3,6 +3,8 @@ from datetime import datetime from ipaddress import IPv4Address from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.protocols.ntp import NTPPacket from primaite.simulator.system.services.service import Service, ServiceOperatingState @@ -12,9 +14,16 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) -class NTPClient(Service): +class NTPClient(Service, identifier="NTPClient"): """Represents a NTP client as a service.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for NTPClient.""" + + type: str = "NTPClient" + + config: "NTPClient.ConfigSchema" = Field(default_factory=lambda: NTPClient.ConfigSchema()) + ntp_server: Optional[IPv4Address] = None "The NTP server the client sends requests to." time: Optional[datetime] = None diff --git a/src/primaite/simulator/system/services/ntp/ntp_server.py b/src/primaite/simulator/system/services/ntp/ntp_server.py index 6e73ccc6..7af33893 100644 --- a/src/primaite/simulator/system/services/ntp/ntp_server.py +++ b/src/primaite/simulator/system/services/ntp/ntp_server.py @@ -2,6 +2,8 @@ from datetime import datetime from typing import Dict, Optional +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.protocols.ntp import NTPPacket from primaite.simulator.system.services.service import Service @@ -11,9 +13,16 @@ from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) -class NTPServer(Service): +class NTPServer(Service, identifier="NTPServer"): """Represents a NTP server as a service.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for NTPServer.""" + + type: str = "NTPServer" + + config: "NTPServer.ConfigSchema" = Field(default_factory=lambda: NTPServer.ConfigSchema()) + def __init__(self, **kwargs): kwargs["name"] = "NTPServer" kwargs["port"] = PORT_LOOKUP["NTP"] diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 3dc080b4..c4a73301 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -1,10 +1,12 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from __future__ import annotations -from abc import abstractmethod +from abc import ABC, abstractmethod from enum import Enum from typing import Any, ClassVar, Dict, Optional, Type +from pydantic import Field + from primaite import getLogger from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestPermissionValidator, RequestType @@ -37,6 +39,13 @@ class Service(IOSoftware): Services are programs that run in the background and may perform input/output operations. """ + class ConfigSchema(IOSoftware.ConfigSchema, ABC): + """Config Schema for Service class.""" + + type: str + + config: "Service.ConfigSchema" = Field(default_factory=lambda: Service.ConfigSchema()) + operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED "The current operating state of the Service." @@ -52,7 +61,7 @@ class Service(IOSoftware): def __init__(self, **kwargs): super().__init__(**kwargs) - def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None: + def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: """ Register a hostnode type. @@ -60,15 +69,30 @@ class Service(IOSoftware): :type identifier: str :raises ValueError: When attempting to register an hostnode with a name that is already allocated. """ - if identifier == "default": + super().__init_subclass__(**kwargs) + if identifier is None: return # Enforce lowercase registry entries because it makes comparisons everywhere else much easier. identifier = identifier.lower() - super().__init_subclass__(**kwargs) if identifier in cls._registry: raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.") cls._registry[identifier] = cls + @classmethod + def from_config(cls, config: Dict) -> "Service": + """Create a service from a config dictionary. + + :param config: dict of options for service components constructor + :type config: dict + :return: The service component. + :rtype: Service + """ + if config["type"] not in cls._registry: + raise ValueError(f"Invalid service type {config['type']}") + service_class = cls._registry[config["type"]] + service_object = service_class(config=service_class.ConfigSchema(**config)) + return service_object + def _can_perform_action(self) -> bool: """ Checks if the service can perform actions. @@ -232,14 +256,14 @@ class Service(IOSoftware): def disable(self) -> bool: """Disable the service.""" - self.sys_log.info(f"Disabling Application {self.name}") + self.sys_log.info(f"Disabling Service {self.name}") self.operating_state = ServiceOperatingState.DISABLED return True def enable(self) -> bool: """Enable the disabled service.""" if self.operating_state == ServiceOperatingState.DISABLED: - self.sys_log.info(f"Enabling Application {self.name}") + self.sys_log.info(f"Enabling Service {self.name}") self.operating_state = ServiceOperatingState.STOPPED return True return False diff --git a/src/primaite/simulator/system/services/terminal/terminal.py b/src/primaite/simulator/system/services/terminal/terminal.py index e26e77f6..bda8bad3 100644 --- a/src/primaite/simulator/system/services/terminal/terminal.py +++ b/src/primaite/simulator/system/services/terminal/terminal.py @@ -7,7 +7,7 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Optional, Union from uuid import uuid4 -from pydantic import BaseModel +from pydantic import BaseModel, Field from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -129,9 +129,16 @@ class RemoteTerminalConnection(TerminalClientConnection): return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id) -class Terminal(Service): +class Terminal(Service, identifier="Terminal"): """Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for Terminal.""" + + type: str = "Terminal" + + config: "Terminal.ConfigSchema" = Field(default_factory=lambda: Terminal.ConfigSchema()) + _client_connection_requests: Dict[str, Optional[Union[str, TerminalClientConnection]]] = {} """Dictionary of connect requests made to remote nodes.""" @@ -179,7 +186,7 @@ class Terminal(Service): return RequestResponse(status="failure", data={}) rm.add_request( - "ssh_to_remote", + "node_session_remote_login", request_type=RequestType(func=_remote_login), ) @@ -279,7 +286,6 @@ class Terminal(Service): :param password: Password for login. :return: boolean, True if successful, else False """ - # TODO: Un-comment this when UserSessionManager is merged. connection_uuid = self.parent.user_session_manager.local_login(username=username, password=password) if connection_uuid: self.sys_log.info(f"{self.name}: Login request authorised, connection uuid: {connection_uuid}") @@ -406,7 +412,6 @@ class Terminal(Service): if isinstance(payload, SSHPacket): if payload.transport_message == SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST: # validate & add connection - # TODO: uncomment this as part of 2781 username = payload.user_account.username password = payload.user_account.password connection_id = self.parent.user_session_manager.remote_login( diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py index 1aab374d..51724002 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/src/primaite/simulator/system/services/web_server/web_server.py @@ -3,6 +3,8 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Optional from urllib.parse import urlparse +from pydantic import Field + from primaite import getLogger from primaite.simulator.network.protocols.http import ( HttpRequestMethod, @@ -19,9 +21,16 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP _LOGGER = getLogger(__name__) -class WebServer(Service): +class WebServer(Service, identifier="WebServer"): """Class used to represent a Web Server Service in simulation.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for WebServer.""" + + type: str = "WebServer" + + config: "WebServer.ConfigSchema" = Field(default_factory=lambda: WebServer.ConfigSchema()) + response_codes_this_timestep: List[HttpStatusCode] = [] def describe_state(self) -> Dict: diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 9a30f3e3..950f77c6 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -1,13 +1,13 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK import copy -from abc import abstractmethod +from abc import ABC, abstractmethod from datetime import datetime from enum import Enum from ipaddress import IPv4Address, IPv4Network from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Union from prettytable import MARKDOWN, PrettyTable -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from primaite.interface.request import RequestResponse from primaite.simulator.core import RequestManager, RequestType, SimComponent @@ -70,7 +70,7 @@ class SoftwareCriticality(Enum): "The highest level of criticality." -class Software(SimComponent): +class Software(SimComponent, ABC): """ A base class representing software in a simulator environment. @@ -78,14 +78,22 @@ class Software(SimComponent): It outlines the fundamental attributes and behaviors expected of any software in the simulation. """ + class ConfigSchema(BaseModel, ABC): + """Configurable options for all software.""" + + model_config = ConfigDict(extra="forbid") + starting_health_state: SoftwareHealthState = SoftwareHealthState.UNUSED + criticality: SoftwareCriticality = SoftwareCriticality.LOWEST + fixing_duration: int = 2 + + config: ConfigSchema = Field(default_factory=lambda: Software.ConfigSchema()) + name: str "The name of the software." health_state_actual: SoftwareHealthState = SoftwareHealthState.UNUSED "The actual health state of the software." health_state_visible: SoftwareHealthState = SoftwareHealthState.UNUSED "The health state of the software visible to the red agent." - criticality: SoftwareCriticality = SoftwareCriticality.LOWEST - "The criticality level of the software." fixing_count: int = 0 "The count of patches applied to the software, defaults to 0." scanning_count: int = 0 @@ -100,11 +108,13 @@ class Software(SimComponent): "The FileSystem of the Node the Software is installed on." folder: Optional[Folder] = None "The folder on the file system the Software uses." - fixing_duration: int = 2 - "The number of ticks it takes to patch the software." _fixing_countdown: Optional[int] = None "Current number of ticks left to patch the software." + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.health_state_actual = self.config.starting_health_state # don't remove this + def _init_request_manager(self) -> RequestManager: """ Initialise the request manager. @@ -152,7 +162,7 @@ class Software(SimComponent): { "health_state_actual": self.health_state_actual.value, "health_state_visible": self.health_state_visible.value, - "criticality": self.criticality.value, + "criticality": self.config.criticality.value, "fixing_count": self.fixing_count, "scanning_count": self.scanning_count, "revealed_to_red": self.revealed_to_red, @@ -201,7 +211,7 @@ class Software(SimComponent): def fix(self) -> bool: """Perform a fix on the software.""" if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD): - self._fixing_countdown = self.fixing_duration + self._fixing_countdown = self.config.fixing_duration self.set_health_state(SoftwareHealthState.FIXING) return True return False @@ -233,7 +243,7 @@ class Software(SimComponent): super().pre_timestep(timestep) -class IOSoftware(Software): +class IOSoftware(Software, ABC): """ Represents software in a simulator environment that is capable of input/output operations. @@ -243,6 +253,13 @@ class IOSoftware(Software): required. """ + class ConfigSchema(Software.ConfigSchema, ABC): + """Configuration options for all IO Software.""" + + listen_on_ports: Set[Port] = Field(default_factory=set) + + config: ConfigSchema = Field(default_factory=lambda: IOSoftware.ConfigSchema()) + installing_count: int = 0 "The number of times the software has been installed. Default is 0." max_sessions: int = 100 @@ -260,6 +277,10 @@ class IOSoftware(Software): _connections: Dict[str, Dict] = {} "Active connections." + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.listen_on_ports = self.config.listen_on_ports + @abstractmethod def describe_state(self) -> Dict: """ diff --git a/src/primaite/utils/session_output_reader.py b/src/primaite/utils/session_output_reader.py index f25bbe6a..f27f6143 100644 --- a/src/primaite/utils/session_output_reader.py +++ b/src/primaite/utils/session_output_reader.py @@ -3,7 +3,6 @@ raise DeprecationWarning( "Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet." ) -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from pathlib import Path from typing import Any, Dict, Tuple, Union diff --git a/src/primaite/utils/session_output_writer.py b/src/primaite/utils/session_output_writer.py index a8cefe35..4049a8c1 100644 --- a/src/primaite/utils/session_output_writer.py +++ b/src/primaite/utils/session_output_writer.py @@ -3,7 +3,6 @@ raise DeprecationWarning( "Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet." ) -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK import csv from logging import Logger from typing import Final, List, Tuple, TYPE_CHECKING, Union diff --git a/src/primaite/utils/validation/ipv4_address.py b/src/primaite/utils/validation/ipv4_address.py index c385ed1e..b2b8b72e 100644 --- a/src/primaite/utils/validation/ipv4_address.py +++ b/src/primaite/utils/validation/ipv4_address.py @@ -31,7 +31,7 @@ def ipv4_validator(v: Any) -> IPv4Address: IPV4Address: Final[Annotated] = Annotated[IPv4Address, BeforeValidator(ipv4_validator)] """ -IPv4Address with with IPv4Address with with pre-validation and auto-conversion from str using ipv4_validator.. +IPv4Address with pre-validation and auto-conversion from str using ipv4_validator.. This type is essentially an IPv4Address from the standard library's ipaddress module, but with added validation logic. If you use this custom type, the ipv4_validator function diff --git a/tests/assets/configs/action_penalty.yaml b/tests/assets/configs/action_penalty.yaml index 1771ba5f..3e57f579 100644 --- a/tests/assets/configs/action_penalty.yaml +++ b/tests/assets/configs/action_penalty.yaml @@ -69,8 +69,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -96,490 +96,423 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_nodename: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_nodename: router_1 + target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index c83cadc8..9cf95a64 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/tests/assets/configs/bad_primaite_session.yaml @@ -13,60 +13,19 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null - action_space: - action_list: - - type: DONOTHING - options: - nodes: - - node_name: client_2 - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_nics_per_node: 2 - max_acl_rules: 10 - - 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 + action_probabilities: + 0: 1.0 - ref: data_manipulation_attacker team: RED type: RedDatabaseCorruptingAgent - - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - - type: NODE_FILE_DELETE - - type: NODE_FILE_CORRUPT - - type: NODE_OS_SCAN - options: - nodes: - - node_name: client_1 - 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 + possible_start_nodes: [client_1,] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -115,8 +74,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -142,381 +101,327 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" + action: "node_file_checkhash" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" + action: "node_folder_checkhash" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 19: # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 20: - action: "NODE_STARTUP" + action: "node_startup" options: - node_id: 5 + node_name: client_1 21: - action: "NODE_RESET" + action: "node_reset" options: - node_id: 5 + node_name: client_1 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 24: # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 25: # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 26: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 27: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 28: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 29: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 30: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 31: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 32: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 33: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 34: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 35: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 36: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 37: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 38: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 39: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 40: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 41: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 42: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 43: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 44: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 45: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 46: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 47: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 48: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 49: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 50: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 51: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 52: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 53: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 - - - - options: - nodes: - - node_name: domain_controller - - node_name: web_server - - node_name: database_server - - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/basic_firewall.yaml b/tests/assets/configs/basic_firewall.yaml index 0253a4d2..3a62c75c 100644 --- a/tests/assets/configs/basic_firewall.yaml +++ b/tests/assets/configs/basic_firewall.yaml @@ -27,39 +27,22 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - 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 + node_name: client_2 + application_name: WebBrowser agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + action_probabilities: + 0: 0.4 + 1: 0.6 simulation: network: diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index 799bb571..a39bf876 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -30,41 +30,26 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - 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 + node_name: client_2 + application_name: WebBrowser reward_function: reward_components: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 - - + action_probabilities: + 0: 0.4 + 1: 0.6 - ref: defender team: BLUE @@ -103,8 +88,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -122,28 +107,10 @@ agents: options: {} action_space: - action_list: - - type: DONOTHING - action_map: 0: - action: DONOTHING + action: do_nothing options: {} - options: - nodes: - - node_name: switch - - node_name: client_1 - - node_name: client_2 - - node_name: client_3 - 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_list: - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.23 reward_function: reward_components: diff --git a/tests/assets/configs/data_manipulation.yaml b/tests/assets/configs/data_manipulation.yaml index bddea1a0..726c9ab0 100644 --- a/tests/assets/configs/data_manipulation.yaml +++ b/tests/assets/configs/data_manipulation.yaml @@ -30,35 +30,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -79,35 +66,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: WebBrowser reward_function: reward_components: @@ -127,34 +101,12 @@ agents: - ref: data_manipulation_attacker team: RED type: RedDatabaseCorruptingAgent - - observation_space: null - - 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 + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -187,7 +139,7 @@ agents: num_files: 1 num_nics: 2 include_num_access: false - include_nmne: true + include_nmne: true monitored_traffic: icmp: - NONE @@ -208,8 +160,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -235,490 +187,423 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/dmz_network.yaml b/tests/assets/configs/dmz_network.yaml index 52316260..b4d018c8 100644 --- a/tests/assets/configs/dmz_network.yaml +++ b/tests/assets/configs/dmz_network.yaml @@ -52,39 +52,22 @@ agents: - ref: client_1_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - options: - nodes: - - node_name: client_1 - 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 + node_name: client_1 + application_name: WebBrowser agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + action_probabilities: + 0: 0.4 + 1: 0.6 simulation: diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index 3d60eb6e..41b7fce9 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/tests/assets/configs/eval_only_primaite_session.yaml @@ -13,72 +13,47 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING action_map: 0: - action: DONOTHING + action: do_nothing options: {} - options: - nodes: - - node_name: client_2 - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_nics_per_node: 2 - max_acl_rules: 10 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 + action_probabilities: + 0: 1.0 - ref: data_manipulation_attacker team: RED type: RedDatabaseCorruptingAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - - type: NODE_FILE_DELETE - - type: NODE_FILE_CORRUPT - - type: NODE_OS_SCAN action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - options: - nodes: - - node_name: client_1 - applications: - - application_name: DataManipulationBot - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 + node_name: client_1 + application_name: DataManipulationBot 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 + possible_start_nodes: [client_1,] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -127,8 +102,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -154,381 +129,327 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" + action: "node_file_checkhash" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 1 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" + action: "node_folder_checkhash" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 1 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 19: # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 20: - action: "NODE_STARTUP" + action: "node_startup" options: - node_id: 5 + node_name: client_1 21: - action: "NODE_RESET" + action: "node_reset" options: - node_id: 5 + node_name: client_1 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 24: # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 25: # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 26: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 27: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 28: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 29: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 30: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 31: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 32: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 33: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 34: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 35: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 36: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 37: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 38: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 39: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 40: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 41: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 42: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 43: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 44: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 45: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 46: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 47: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 48: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 49: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 50: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 51: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 52: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 53: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 - - - - options: - nodes: - - node_name: domain_controller - - node_name: web_server - - node_name: database_server - - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/extended_config.yaml b/tests/assets/configs/extended_config.yaml index e1a06938..bff58ebd 100644 --- a/tests/assets/configs/extended_config.yaml +++ b/tests/assets/configs/extended_config.yaml @@ -30,35 +30,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -79,35 +66,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: DatabaseClient reward_function: reward_components: @@ -128,33 +102,12 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - 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 + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -208,8 +161,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -235,491 +188,426 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_file_checkhash" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_checkhash" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - - reward_function: reward_components: - type: DATABASE_FILE_INTEGRITY diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index 2292616d..4b11dbcc 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -77,8 +77,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -95,181 +95,156 @@ agents: options: {} action_space: - action_list: - - type: DONOTHING - - type: FIREWALL_ACL_ADDRULE - - type: FIREWALL_ACL_REMOVERULE - - type: NETWORK_PORT_DISABLE - - type: NETWORK_PORT_ENABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: + type: firewall_acl_add_rule target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: inbound position: 1 - permission: 1 - source_ip_id: 2 # client 1 - dest_ip_id: 1 # ALL - source_port_id: 1 - dest_port_id: 1 - protocol_id: 1 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: PERMIT + src_ip: 192.168.0.10 + dst_ip: ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 2: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: inbound position: 1 3: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: outbound position: 1 - permission: 2 - source_ip_id: 2 # client 1 - dest_ip_id: 1 # ALL - source_port_id: 2 - dest_port_id: 3 - protocol_id: 2 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.0.10 # client 1 + dst_ip: ALL + src_port: ARP + dst_port: DNS + protocol_name: ICMP + src_wildcard: NONE + dst_wildcard: NONE 4: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: outbound position: 1 5: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: inbound position: 1 - permission: 2 - source_ip_id: 3 # dmz_server - dest_ip_id: 2 # client_1 - source_port_id: 4 - dest_port_id: 4 - protocol_id: 4 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.10 # dmz_server + dst_ip: 192.168.0.10 # client_1 + src_port: HTTP + dst_port: HTTP + protocol_name: UDP + src_wildcard: NONE + dst_wildcard: NONE 6: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: inbound position: 1 7: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: outbound position: 2 - permission: 2 - source_ip_id: 3 # dmz_server - dest_ip_id: 2 # client_1 - source_port_id: 4 - dest_port_id: 4 - protocol_id: 3 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.10 # dmz_server + dst_ip: 192.168.0.10 # client_1 + src_port: HTTP + dst_port: HTTP + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 8: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: outbound position: 2 9: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: inbound position: 10 - permission: 2 - source_ip_id: 4 # external_computer - dest_ip_id: 3 # dmz - source_port_id: 5 - dest_port_id: 5 - protocol_id: 2 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.20.10 # external_computer + dst_ip: 192.168.10.10 # dmz + src_port: POSTGRES_SERVER + dst_port: POSTGRES_SERVER + protocol_name: ICMP + src_wildcard: NONE + dst_wildcard: NONE 10: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: inbound position: 10 11: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: outbound position: 1 - permission: 2 - source_ip_id: 4 # external_computer - dest_ip_id: 2 # client_1 - source_port_id: 1 - dest_port_id: 1 - protocol_id: 1 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.20.10 # external_computer + dst_ip: 192.168.0.10 # client_1 + src_port: NONE + dst_port: NONE + protocol_name: none + src_wildcard: NONE + dst_wildcard: NONE 12: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: outbound position: 1 13: - action: NETWORK_PORT_DISABLE + action: network_port_disable options: + type: network_port_disable target_nodename: firewall - port_id: 3 + port_num: 3 14: - action: NETWORK_PORT_ENABLE + action: network_port_enable options: + type: network_port_enable target_nodename: firewall - port_id: 3 - options: - nodes: - - node_name: client_1 - - node_name: dmz_server - - node_name: external_computer - ip_list: - - 192.168.0.10 - - 192.168.10.10 - - 192.168.20.10 - max_folders_per_node: 2 - max_files_per_folder: 2 - max_services_per_node: 2 - max_nics_per_node: 8 - max_acl_rules: 10 - reward_function: - reward_components: - - type: DUMMY + port_num: 3 - agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 diff --git a/tests/assets/configs/fix_duration_one_item.yaml b/tests/assets/configs/fixing_duration_one_item.yaml similarity index 78% rename from tests/assets/configs/fix_duration_one_item.yaml rename to tests/assets/configs/fixing_duration_one_item.yaml index c74590a1..da5a9993 100644 --- a/tests/assets/configs/fix_duration_one_item.yaml +++ b/tests/assets/configs/fixing_duration_one_item.yaml @@ -27,40 +27,21 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - 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 - + node_name: client_1 + application_name: WebBrowser agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 - + action_probabilities: + 0: 0.4 + 1: 0.6 - ref: defender @@ -100,8 +81,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -119,28 +100,10 @@ agents: options: {} action_space: - action_list: - - type: DONOTHING - action_map: 0: - action: DONOTHING + action: do_nothing options: {} - options: - nodes: - - node_name: switch - - node_name: client_1 - - node_name: client_2 - - node_name: client_3 - 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_list: - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.23 reward_function: reward_components: @@ -185,7 +148,7 @@ simulation: options: db_server_ip: 192.168.1.10 server_password: arcd - fix_duration: 1 + fixing_duration: 1 - type: DataManipulationBot options: port_scan_p_of_success: 0.8 @@ -206,7 +169,7 @@ simulation: arcd.com: 192.168.1.10 - type: DatabaseService options: - fix_duration: 5 + fixing_duration: 5 backup_server_ip: 192.168.1.10 - type: WebServer - type: FTPClient diff --git a/tests/assets/configs/install_and_configure_apps.yaml b/tests/assets/configs/install_and_configure_apps.yaml index 6b548f7e..2baca409 100644 --- a/tests/assets/configs/install_and_configure_apps.yaml +++ b/tests/assets/configs/install_and_configure_apps.yaml @@ -18,80 +18,63 @@ agents: team: BLUE type: ProxyAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_INSTALL - - type: CONFIGURE_DATABASE_CLIENT - - type: CONFIGURE_DOSBOT - - type: CONFIGURE_RANSOMWARE_SCRIPT - - type: NODE_APPLICATION_REMOVE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_INSTALL + action: node_application_install options: - node_id: 0 + node_name: client_1 application_name: DatabaseClient 2: - action: NODE_APPLICATION_INSTALL + action: node_application_install options: - node_id: 1 + node_name: client_2 application_name: RansomwareScript 3: - action: NODE_APPLICATION_INSTALL + action: node_application_install options: - node_id: 2 + node_name: client_3 application_name: DoSBot 4: - action: CONFIGURE_DATABASE_CLIENT + action: configure_database_client options: - node_id: 0 - config: - server_ip_address: 10.0.0.5 + node_name: client_1 + server_ip_address: 10.0.0.5 5: - action: CONFIGURE_DATABASE_CLIENT + action: configure_database_client options: - node_id: 0 - config: - server_password: correct_password + node_name: client_1 + server_password: correct_password 6: - action: CONFIGURE_RANSOMWARE_SCRIPT + action: configure_ransomware_script options: - node_id: 1 - config: - server_ip_address: 10.0.0.5 - server_password: correct_password - payload: ENCRYPT + node_name: client_2 + server_ip_address: 10.0.0.5 + server_password: correct_password + payload: ENCRYPT 7: - action: CONFIGURE_DOSBOT + action: configure_dos_bot options: - node_id: 2 - config: - target_ip_address: 10.0.0.5 - target_port: POSTGRES_SERVER - payload: DELETE - repeat: true - port_scan_p_of_success: 1.0 - dos_intensity: 1.0 - max_sessions: 1000 + node_name: client_3 + target_ip_address: 10.0.0.5 + target_port: POSTGRES_SERVER + payload: DELETE + repeat: true + port_scan_p_of_success: 1.0 + dos_intensity: 1.0 + max_sessions: 1000 8: - action: NODE_APPLICATION_INSTALL + action: node_application_install options: - node_id: 1 + node_name: client_2 application_name: DatabaseClient - options: - nodes: - - node_name: client_1 - - node_name: client_2 - - node_name: client_3 - ip_list: [] - reward_function: - reward_components: - - type: DUMMY + agent_settings: + flatten_obs: True + action_masking: False simulation: network: diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index a2d64605..93baf4af 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -26,35 +26,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -75,35 +62,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: WebBrowser reward_function: reward_components: @@ -124,36 +98,12 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - - type: NODE_FILE_DELETE - - type: NODE_FILE_CORRUPT - - type: NODE_OS_SCAN - 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 + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender_1 team: BLUE @@ -202,8 +152,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -229,489 +179,423 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - + node_name: client_2 + nic_num: 1 reward_function: reward_components: @@ -782,8 +666,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -809,494 +693,425 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - options: - target_router: router_1 - - type: ROUTER_ACL_REMOVERULE - options: - target_router: router_1 - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 - - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - reward_function: reward_components: - type: DATABASE_FILE_INTEGRITY diff --git a/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml b/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml index c5508f13..c2f79144 100644 --- a/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml +++ b/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml @@ -22,32 +22,17 @@ agents: - ref: client_1_red_nmap team: RED type: ProbabilisticAgent - observation_space: null + action_space: - options: - nodes: - - node_name: client_1 - applications: - - application_name: NMAP - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_applications_per_node: 1 - action_list: - - type: NODE_NMAP_NETWORK_SERVICE_RECON action_map: 0: - action: NODE_NMAP_NETWORK_SERVICE_RECON + action: node_network_service_recon options: source_node: client_1 target_ip_address: 192.168.10.0/24 target_port: 80 target_protocol: tcp - - reward_function: - reward_components: - - type: DUMMY - + show: false agent_settings: action_probabilities: 0: 1.0 diff --git a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml index 33ba3d19..f6d549e8 100644 --- a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml +++ b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml @@ -22,29 +22,15 @@ agents: - ref: client_1_red_nmap team: RED type: ProbabilisticAgent - observation_space: null + action_space: - options: - nodes: - - node_name: client_1 - applications: - - application_name: NMAP - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_applications_per_node: 1 - action_list: - - type: NODE_NMAP_PING_SCAN action_map: 0: - action: NODE_NMAP_PING_SCAN + action: node_nmap_ping_scan options: source_node: client_1 target_ip_address: 192.168.1.0/24 - - reward_function: - reward_components: - - type: DUMMY + show: False agent_settings: action_probabilities: diff --git a/tests/assets/configs/nmap_port_scan_red_agent_config.yaml b/tests/assets/configs/nmap_port_scan_red_agent_config.yaml index 8ed715c1..873401b9 100644 --- a/tests/assets/configs/nmap_port_scan_red_agent_config.yaml +++ b/tests/assets/configs/nmap_port_scan_red_agent_config.yaml @@ -22,22 +22,11 @@ agents: - ref: client_1_red_nmap team: RED type: ProbabilisticAgent - observation_space: null + action_space: - options: - nodes: - - node_name: client_1 - applications: - - application_name: NMAP - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_applications_per_node: 1 - action_list: - - type: NODE_NMAP_PORT_SCAN action_map: 0: - action: NODE_NMAP_PORT_SCAN + action: node_nmap_port_scan options: source_node: client_1 target_ip_address: 192.168.10.0/24 @@ -47,10 +36,7 @@ agents: - 80 - 123 - 219 - - reward_function: - reward_components: - - type: DUMMY + show: false agent_settings: action_probabilities: diff --git a/tests/assets/configs/scenario_with_placeholders/greens_1.yaml b/tests/assets/configs/scenario_with_placeholders/greens_1.yaml index 98d2392a..677cd5a5 100644 --- a/tests/assets/configs/scenario_with_placeholders/greens_1.yaml +++ b/tests/assets/configs/scenario_with_placeholders/greens_1.yaml @@ -6,25 +6,17 @@ agents: &greens action_probabilities: 0: 0.2 1: 0.8 - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DatabaseClient action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client + application_name: DatabaseClient reward_function: reward_components: diff --git a/tests/assets/configs/scenario_with_placeholders/greens_2.yaml b/tests/assets/configs/scenario_with_placeholders/greens_2.yaml index 17a5977b..eb7823f8 100644 --- a/tests/assets/configs/scenario_with_placeholders/greens_2.yaml +++ b/tests/assets/configs/scenario_with_placeholders/greens_2.yaml @@ -6,25 +6,17 @@ agents: &greens action_probabilities: 0: 0.95 1: 0.05 - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DatabaseClient action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client + application_name: DatabaseClient reward_function: reward_components: diff --git a/tests/assets/configs/scenario_with_placeholders/reds_1.yaml b/tests/assets/configs/scenario_with_placeholders/reds_1.yaml index 31675a0b..0170143f 100644 --- a/tests/assets/configs/scenario_with_placeholders/reds_1.yaml +++ b/tests/assets/configs/scenario_with_placeholders/reds_1.yaml @@ -3,24 +3,9 @@ reds: &reds team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DataManipulationBot - - reward_function: - reward_components: - - type: DUMMY - agent_settings: - start_settings: - start_step: 10 - frequency: 10 - variance: 0 + possible_start_nodes: [client,] + target_application: DataManipulationBot + start_step: 10 + frequency: 10 + variance: 0 diff --git a/tests/assets/configs/scenario_with_placeholders/reds_2.yaml b/tests/assets/configs/scenario_with_placeholders/reds_2.yaml index c5572b89..e14eaa43 100644 --- a/tests/assets/configs/scenario_with_placeholders/reds_2.yaml +++ b/tests/assets/configs/scenario_with_placeholders/reds_2.yaml @@ -2,25 +2,9 @@ reds: &reds - ref: red_B team: RED type: RedDatabaseCorruptingAgent - - observation_space: null - - action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - options: - nodes: - - node_name: client - applications: - - application_name: DataManipulationBot - - reward_function: - reward_components: - - type: DUMMY - agent_settings: - start_settings: - start_step: 3 - frequency: 2 - variance: 1 + possible_start_nodes: [client_1,] + target_application: DataManipulationBot + start_step: 3 + frequency: 2 + variance: 1 diff --git a/tests/assets/configs/scenario_with_placeholders/scenario.yaml b/tests/assets/configs/scenario_with_placeholders/scenario.yaml index ef930a1a..7ea0145a 100644 --- a/tests/assets/configs/scenario_with_placeholders/scenario.yaml +++ b/tests/assets/configs/scenario_with_placeholders/scenario.yaml @@ -54,65 +54,46 @@ agents: - server:eth-1<->switch_1:eth-2 action_space: - action_list: - - type: DONOTHING - - type: NODE_SHUTDOWN - - type: NODE_STARTUP - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 0 + node_name: client 2: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 1 + node_name: server 3: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: client 4: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: server 5: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 0 - nic_id: 0 + node_name: client + nic_num: 1 6: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 1 - nic_id: 0 + node_name: server + nic_num: 1 7: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 0 - nic_id: 0 + node_name: client + nic_num: 1 8: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 1 - nic_id: 0 - options: - nodes: - - node_name: client - - node_name: server - - max_folders_per_node: 0 - max_files_per_folder: 0 - max_services_per_node: 0 - max_nics_per_node: 1 - max_acl_rules: 0 - ip_list: - - 192.168.1.2 - - 192.168.1.3 + node_name: server + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/shared_rewards.yaml b/tests/assets/configs/shared_rewards.yaml index 81cb85f7..d5615a72 100644 --- a/tests/assets/configs/shared_rewards.yaml +++ b/tests/assets/configs/shared_rewards.yaml @@ -29,35 +29,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -78,35 +65,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: DatabaseClient reward_function: reward_components: @@ -123,33 +97,12 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - 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 + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -198,8 +151,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -225,490 +178,423 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" + action: "node_file_checkhash" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" + action: "node_folder_checkhash" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/software_fix_duration.yaml b/tests/assets/configs/software_fixing_duration.yaml similarity index 75% rename from tests/assets/configs/software_fix_duration.yaml rename to tests/assets/configs/software_fixing_duration.yaml index 6a705b37..f685b420 100644 --- a/tests/assets/configs/software_fix_duration.yaml +++ b/tests/assets/configs/software_fixing_duration.yaml @@ -27,40 +27,21 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - 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 - + node_name: client_2 + application_name: WebBrowser agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 - + action_probabilities: + 0: 0.4 + 1: 0.6 - ref: defender @@ -100,8 +81,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -119,28 +100,10 @@ agents: options: {} action_space: - action_list: - - type: DONOTHING - action_map: 0: - action: DONOTHING + action: do_nothing options: {} - options: - nodes: - - node_name: switch - - node_name: client_1 - - node_name: client_2 - - node_name: client_3 - 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_list: - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.23 reward_function: reward_components: @@ -179,19 +142,19 @@ simulation: applications: - type: NMAP options: - fix_duration: 1 + fixing_duration: 1 - type: RansomwareScript options: - fix_duration: 1 + fixing_duration: 1 - type: WebBrowser options: target_url: http://arcd.com/users/ - fix_duration: 1 + fixing_duration: 1 - type: DatabaseClient options: db_server_ip: 192.168.1.10 server_password: arcd - fix_duration: 1 + fixing_duration: 1 - type: DataManipulationBot options: port_scan_p_of_success: 0.8 @@ -199,43 +162,44 @@ simulation: payload: "DELETE" server_ip: 192.168.1.21 server_password: arcd - fix_duration: 1 + fixing_duration: 1 - type: DoSBot options: target_ip_address: 192.168.10.21 payload: SPOOF DATA port_scan_p_of_success: 0.8 - fix_duration: 1 + fixing_duration: 1 services: - type: DNSClient options: - fix_duration: 3 + dns_server: 192.168.1.10 + fixing_duration: 3 - type: DNSServer options: - fix_duration: 3 + fixing_duration: 3 domain_mapping: arcd.com: 192.168.1.10 - type: DatabaseService options: backup_server_ip: 192.168.1.10 - fix_duration: 3 + fixing_duration: 3 - type: WebServer options: - fix_duration: 3 + fixing_duration: 3 - type: FTPClient options: - fix_duration: 3 + fixing_duration: 3 - type: FTPServer options: server_password: arcd - fix_duration: 3 + fixing_duration: 3 - type: NTPClient options: ntp_server_ip: 192.168.1.10 - fix_duration: 3 + fixing_duration: 3 - type: NTPServer options: - fix_duration: 3 + fixing_duration: 3 - hostname: client_2 type: computer ip_address: 192.168.10.22 diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index 3a3a6890..25bc38e6 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/tests/assets/configs/test_application_install.yaml @@ -29,35 +29,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_2 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_2 + application_name: DatabaseClient reward_function: reward_components: @@ -78,35 +65,22 @@ agents: 0: 0.3 1: 0.6 2: 0.1 - observation_space: null + 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 + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: client_1 + application_name: WebBrowser 2: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 1 + node_name: client_1 + application_name: WebBrowser reward_function: reward_components: @@ -127,33 +101,12 @@ agents: team: RED type: RedDatabaseCorruptingAgent - observation_space: null - - 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 + possible_start_nodes: [client_1, client_2] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -202,8 +155,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -229,524 +182,449 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - - type: NODE_APPLICATION_INSTALL - - type: NODE_APPLICATION_REMOVE - - type: NODE_APPLICATION_EXECUTE - - type: CONFIGURE_DOSBOT - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + 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 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 0 + node_name: domain_controller 19: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 0 + node_name: domain_controller 20: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: domain_controller 21: - action: NODE_RESET + action: node_reset options: - node_id: 0 + node_name: domain_controller 22: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 1 + node_name: web_server 23: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 1 + node_name: web_server 24: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: web_server 25: - action: NODE_RESET + action: node_reset options: - node_id: 1 + node_name: web_server 26: # old action num: 18 - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 27: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 2 + node_name: database_server 28: - action: NODE_STARTUP + action: node_startup options: - node_id: 2 + node_name: database_server 29: - action: NODE_RESET + action: node_reset options: - node_id: 2 + node_name: database_server 30: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 3 + node_name: backup_server 31: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 3 + node_name: backup_server 32: - action: NODE_STARTUP + action: node_startup options: - node_id: 3 + node_name: backup_server 33: - action: NODE_RESET + action: node_reset options: - node_id: 3 + node_name: backup_server 34: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 4 + node_name: security_suite 35: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 4 + node_name: security_suite 36: - action: NODE_STARTUP + action: node_startup options: - node_id: 4 + node_name: security_suite 37: - action: NODE_RESET + action: node_reset options: - node_id: 4 + node_name: security_suite 38: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 5 + node_name: client_1 39: # old action num: 19 # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 40: # old action num: 20 - action: NODE_STARTUP + action: node_startup options: - node_id: 5 + node_name: client_1 41: # old action num: 21 - action: NODE_RESET + action: node_reset options: - node_id: 5 + node_name: client_1 42: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 6 + node_name: client_2 43: - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 6 + node_name: client_2 44: - action: NODE_STARTUP + action: node_startup options: - node_id: 6 + node_name: client_2 45: - action: NODE_RESET + action: node_reset options: - node_id: 6 + node_name: client_2 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 48: # old action num: 24 # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 49: # old action num: 25 # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 50: # old action num: 26 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 51: # old action num: 27 - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: - target_router_hostname: router_1 + target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 52: # old action num: 28 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 63: # old action num: 39 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 64: # old action num: 40 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 65: # old action num: 41 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 66: # old action num: 42 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 67: # old action num: 43 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 68: # old action num: 44 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 69: # old action num: 45 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 70: # old action num: 46 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 71: # old action num: 47 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 72: # old action num: 48 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 73: # old action num: 49 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 74: # old action num: 50 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 75: # old action num: 51 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 76: # old action num: 52 - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 77: # old action num: 53 - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 78: - action: NODE_APPLICATION_INSTALL + action: node_application_install options: - node_id: 0 + node_name: domain_controller application_name: DoSBot 79: - action: NODE_APPLICATION_REMOVE + action: node_application_remove options: - node_id: 0 + node_name: domain_controller application_name: DoSBot 80: - action: NODE_APPLICATION_REMOVE + action: node_application_remove options: - node_id: 0 + node_name: domain_controller application_name: WebBrowser 81: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 + node_name: domain_controller + application_name: DoSBot 82: - action: CONFIGURE_DOSBOT + action: configure_dos_bot options: - node_id: 0 - config: - target_ip_address: 192.168.1.14 - target_port: POSTGRES_SERVER - - - - - options: - nodes: - - node_name: domain_controller - applications: - - application_name: DoSBot - - 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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 - + node_name: domain_controller + target_ip_address: 192.168.1.14 + target_port: POSTGRES_SERVER reward_function: reward_components: diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index 27cfa240..2d124981 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/tests/assets/configs/test_primaite_session.yaml @@ -21,73 +21,40 @@ agents: - ref: client_2_green_user team: GREEN type: ProbabilisticAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING action_map: 0: - action: DONOTHING + action: do_nothing options: {} - options: - nodes: - - node_name: client_2 - max_folders_per_node: 1 - max_files_per_folder: 1 - max_services_per_node: 1 - max_nics_per_node: 2 - max_acl_rules: 10 - - 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 + action_probabilities: + 0: 1.0 - ref: data_manipulation_attacker team: RED type: RedDatabaseCorruptingAgent - observation_space: null + action_space: - action_list: - - type: DONOTHING - - type: NODE_APPLICATION_EXECUTE - - type: NODE_FILE_DELETE - - type: NODE_FILE_CORRUPT - - type: NODE_OS_SCAN action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_APPLICATION_EXECUTE + action: node_application_execute options: - node_id: 0 - application_id: 0 - options: - nodes: - - node_name: client_1 - 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 + node_name: client_1 + application_name: DataManipulationBot agent_settings: # options specific to this particular agent type, basically args of __init__(self) - start_settings: - start_step: 25 - frequency: 20 - variance: 5 + possible_start_nodes: [client_1,] + target_application: DataManipulationBot + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE @@ -136,8 +103,8 @@ agents: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -163,391 +130,326 @@ agents: options: {} 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_FIX - - 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: ROUTER_ACL_ADDRULE - - type: ROUTER_ACL_REMOVERULE - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE - action_map: 0: - action: DONOTHING + action: do_nothing options: {} # scan webapp service 1: - action: NODE_SERVICE_SCAN + action: node_service_scan options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # stop webapp service 2: - action: NODE_SERVICE_STOP + action: node_service_stop options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer # start webapp service 3: - action: "NODE_SERVICE_START" + action: "node_service_start" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 4: - action: "NODE_SERVICE_PAUSE" + action: "node_service_pause" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 5: - action: "NODE_SERVICE_RESUME" + action: "node_service_resume" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 6: - action: "NODE_SERVICE_RESTART" + action: "node_service_restart" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 7: - action: "NODE_SERVICE_DISABLE" + action: "node_service_disable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 8: - action: "NODE_SERVICE_ENABLE" + action: "node_service_enable" options: - node_id: 1 - service_id: 0 + node_name: web_server + service_name: WebServer 9: # check database.db file - action: "NODE_FILE_SCAN" + action: "node_file_scan" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 10: - action: "NODE_FILE_CHECKHASH" + action: "node_file_checkhash" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 11: - action: "NODE_FILE_DELETE" + action: "node_file_delete" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 12: - action: "NODE_FILE_REPAIR" + action: "node_file_repair" options: - node_id: 2 - folder_id: 0 - file_id: 0 + node_name: database_server + folder_name: database + file_name: database.db 13: - action: "NODE_SERVICE_FIX" + action: "node_service_fix" options: - node_id: 2 - service_id: 0 + node_name: database_server + service_name: DatabaseService 14: - action: "NODE_FOLDER_SCAN" + action: "node_folder_scan" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 15: - action: "NODE_FOLDER_CHECKHASH" + action: "node_folder_checkhash" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 16: - action: "NODE_FOLDER_REPAIR" + action: "node_folder_repair" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 17: - action: "NODE_FOLDER_RESTORE" + action: "node_folder_restore" options: - node_id: 2 - folder_id: 0 + node_name: database_server + folder_name: database 18: - action: "NODE_OS_SCAN" + action: "node_os_scan" options: - node_id: 2 + node_name: database_server 19: # shutdown client 1 - action: "NODE_SHUTDOWN" + action: "node_shutdown" options: - node_id: 5 + node_name: client_1 20: - action: "NODE_STARTUP" + action: "node_startup" options: - node_id: 5 + node_name: client_1 21: - action: "NODE_RESET" + action: "node_reset" options: - node_id: 5 + node_name: client_1 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: ALL # ALL + src_port: ALL + dst_port: ALL + protocol_name: ALL + src_wildcard: NONE + dst_wildcard: NONE 24: # block tcp traffic from client 1 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 25: # block tcp traffic from client 2 to web app - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.12 # web server + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 26: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.21 # client 1 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 27: - action: "ROUTER_ACL_ADDRULE" + action: "router_acl_add_rule" options: target_router: router_1 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 - source_wildcard_id: 0 - dest_wildcard_id: 0 + permission: DENY + src_ip: 192.168.10.22 # client 2 + dst_ip: 192.168.1.14 # database + src_port: ALL + dst_port: ALL + protocol_name: TCP + src_wildcard: NONE + dst_wildcard: NONE 28: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 0 29: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 1 30: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 2 31: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 3 32: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 4 33: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 5 34: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 6 35: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 7 36: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 8 37: - action: "ROUTER_ACL_REMOVERULE" + action: "router_acl_remove_rule" options: target_router: router_1 position: 9 38: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 39: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 0 - nic_id: 0 + node_name: domain_controller + nic_num: 1 40: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 41: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 1 - nic_id: 0 + node_name: web_server + nic_num: 1 42: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 43: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 2 - nic_id: 0 + node_name: database_server + nic_num: 1 44: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 45: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 3 - nic_id: 0 + node_name: backup_server + nic_num: 1 46: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 47: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 0 + node_name: security_suite + nic_num: 1 48: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 49: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 4 - nic_id: 1 + node_name: security_suite + nic_num: 2 50: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 51: - action: "HOST_NIC_ENABLE" + action: "host_nic_enable" options: - node_id: 5 - nic_id: 0 + node_name: client_1 + nic_num: 1 52: - action: "HOST_NIC_DISABLE" + action: "host_nic_disable" options: - node_id: 6 - nic_id: 0 + node_name: client_2 + nic_num: 1 53: - action: "HOST_NIC_ENABLE" + action: "host_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_list: - - 192.168.1.10 - - 192.168.1.12 - - 192.168.1.14 - - 192.168.1.16 - - 192.168.1.110 - - 192.168.10.21 - - 192.168.10.22 - - 192.168.10.110 + node_name: client_2 + nic_num: 1 reward_function: reward_components: diff --git a/tests/conftest.py b/tests/conftest.py index 08c16537..287f216c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,9 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -from typing import Any, Dict, Tuple +from typing import Any, Dict, Optional, Tuple import pytest import yaml +from pydantic import Field from ray import init as rayinit from primaite import getLogger, PRIMAITE_PATHS @@ -10,6 +11,7 @@ from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent from primaite.game.agent.observations.observation_manager import NestedObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction +from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent from primaite.game.game import PrimaiteGame from primaite.simulator.file_system.file_system import FileSystem from primaite.simulator.network.container import Network @@ -37,9 +39,16 @@ ACTION_SPACE_NODE_ACTION_VALUES = 1 _LOGGER = getLogger(__name__) -class DummyService(Service): +class DummyService(Service, identifier="DummyService"): """Test Service class""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for DummyService.""" + + type: str = "DummyService" + + config: "DummyService.ConfigSchema" = Field(default_factory=lambda: DummyService.ConfigSchema()) + def describe_state(self) -> Dict: return super().describe_state() @@ -56,6 +65,13 @@ class DummyService(Service): class DummyApplication(Application, identifier="DummyApplication"): """Test Application class""" + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for DummyApplication.""" + + type: str = "DummyApplication" + + config: "DummyApplication.ConfigSchema" = Field(default_factory=lambda: DummyApplication.ConfigSchema()) + def __init__(self, **kwargs): kwargs["name"] = "DummyApplication" kwargs["port"] = PORT_LOOKUP["HTTP"] @@ -196,7 +212,7 @@ def example_network() -> Network: # Router 1 - router_1_cfg = {"hostname":"router_1", "type":"router"} + router_1_cfg = {"hostname": "router_1", "type": "router"} # router_1 = Router(hostname="router_1", start_up_duration=0) router_1 = Router.from_config(config=router_1_cfg) @@ -226,27 +242,30 @@ def example_network() -> Network: # # Client 1 - client_1_cfg = {"type": "computer", - "hostname": "client_1", - "ip_address": "192.168.10.21", - "subnet_mask": "255.255.255.0", - "default_gateway": "192.168.10.1", - "start_up_duration": 0} + client_1_cfg = { + "type": "computer", + "hostname": "client_1", + "ip_address": "192.168.10.21", + "subnet_mask": "255.255.255.0", + "default_gateway": "192.168.10.1", + "start_up_duration": 0, + } - client_1=Computer.from_config(config=client_1_cfg) + client_1 = Computer.from_config(config=client_1_cfg) client_1.power_on() network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1]) # # Client 2 - client_2_cfg = {"type": "computer", - "hostname": "client_2", - "ip_address": "192.168.10.22", - "subnet_mask": "255.255.255.0", - "default_gateway": "192.168.10.1", - "start_up_duration": 0, - } + client_2_cfg = { + "type": "computer", + "hostname": "client_2", + "ip_address": "192.168.10.22", + "subnet_mask": "255.255.255.0", + "default_gateway": "192.168.10.1", + "start_up_duration": 0, + } client_2 = Computer.from_config(config=client_2_cfg) @@ -255,13 +274,14 @@ def example_network() -> Network: # # Server 1 - server_1_cfg = {"type": "server", - "hostname": "server_1", - "ip_address":"192.168.1.10", - "subnet_mask":"255.255.255.0", - "default_gateway":"192.168.1.1", - "start_up_duration":0, - } + server_1_cfg = { + "type": "server", + "hostname": "server_1", + "ip_address": "192.168.1.10", + "subnet_mask": "255.255.255.0", + "default_gateway": "192.168.1.1", + "start_up_duration": 0, + } server_1 = Server.from_config(config=server_1_cfg) @@ -270,12 +290,13 @@ def example_network() -> Network: # # DServer 2 - server_2_cfg = {"type": "server", - "hostname": "server_2", - "ip_address":"192.168.1.14", - "subnet_mask":"255.255.255.0", - "default_gateway":"192.168.1.1", - "start_up_duration":0, + server_2_cfg = { + "type": "server", + "hostname": "server_2", + "ip_address": "192.168.1.14", + "subnet_mask": "255.255.255.0", + "default_gateway": "192.168.1.1", + "start_up_duration": 0, } server_2 = Server.from_config(config=server_2_cfg) @@ -292,23 +313,16 @@ def example_network() -> Network: return network -class ControlledAgent(AbstractAgent): +class ControlledAgent(AbstractAgent, identifier="ControlledAgent"): """Agent that can be controlled by the tests.""" - def __init__( - self, - agent_name: str, - action_space: ActionManager, - observation_space: ObservationManager, - reward_function: RewardFunction, - ) -> None: - super().__init__( - agent_name=agent_name, - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - ) - self.most_recent_action: Tuple[str, Dict] + config: "ControlledAgent.ConfigSchema" = Field(default_factory=lambda: ControlledAgent.ConfigSchema()) + most_recent_action: Optional[Tuple[str, Dict]] = None + + class ConfigSchema(AbstractAgent.ConfigSchema): + """Configuration Schema for Abstract Agent used in tests.""" + + type: str = "ControlledAgent" def get_action(self, obs: None, timestep: int = 0) -> Tuple[str, Dict]: """Return the agent's most recent action, formatted in CAOS format.""" @@ -403,6 +417,7 @@ def install_stuff_to_sim(sim: Simulation): # 5: Assert that the simulation starts off in the state that we expect assert len(sim.network.nodes) == 6 assert len(sim.network.links) == 5 + # 5.1: Assert the router is correctly configured r = sim.network.router_nodes[0] for i, acl_rule in enumerate(r.acl.acl): @@ -446,97 +461,13 @@ def game_and_agent(): sim = game.simulation install_stuff_to_sim(sim) - actions = [ - {"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_FIX"}, - {"type": "NODE_APPLICATION_EXECUTE"}, - {"type": "NODE_APPLICATION_SCAN"}, - {"type": "NODE_APPLICATION_CLOSE"}, - {"type": "NODE_APPLICATION_FIX"}, - {"type": "NODE_APPLICATION_INSTALL"}, - {"type": "NODE_APPLICATION_REMOVE"}, - {"type": "NODE_FILE_CREATE"}, - {"type": "NODE_FILE_SCAN"}, - {"type": "NODE_FILE_CHECKHASH"}, - {"type": "NODE_FILE_DELETE"}, - {"type": "NODE_FILE_REPAIR"}, - {"type": "NODE_FILE_RESTORE"}, - {"type": "NODE_FILE_CORRUPT"}, - {"type": "NODE_FILE_ACCESS"}, - {"type": "NODE_FOLDER_CREATE"}, - {"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": "ROUTER_ACL_ADDRULE"}, - {"type": "ROUTER_ACL_REMOVERULE"}, - {"type": "HOST_NIC_ENABLE"}, - {"type": "HOST_NIC_DISABLE"}, - {"type": "NETWORK_PORT_ENABLE"}, - {"type": "NETWORK_PORT_DISABLE"}, - {"type": "CONFIGURE_C2_BEACON"}, - {"type": "C2_SERVER_RANSOMWARE_LAUNCH"}, - {"type": "C2_SERVER_RANSOMWARE_CONFIGURE"}, - {"type": "C2_SERVER_TERMINAL_COMMAND"}, - {"type": "C2_SERVER_DATA_EXFILTRATE"}, - {"type": "NODE_ACCOUNTS_CHANGE_PASSWORD"}, - {"type": "SSH_TO_REMOTE"}, - {"type": "SESSIONS_REMOTE_LOGOFF"}, - {"type": "NODE_SEND_REMOTE_COMMAND"}, - ] + config = { + "type": "ControlledAgent", + "ref": "test_agent", + "team": "BLUE", + } - action_space = ActionManager( - actions=actions, # ALL POSSIBLE ACTIONS - nodes=[ - { - "node_name": "client_1", - "applications": [ - {"application_name": "WebBrowser"}, - {"application_name": "DoSBot"}, - {"application_name": "C2Server"}, - ], - "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}], - }, - { - "node_name": "server_1", - "services": [{"service_name": "DNSServer"}], - "applications": [{"application_name": "C2Beacon"}], - }, - {"node_name": "server_2", "services": [{"service_name": "WebServer"}]}, - {"node_name": "router"}, - ], - max_folders_per_node=2, - max_files_per_folder=2, - max_services_per_node=2, - max_applications_per_node=3, - max_nics_per_node=2, - max_acl_rules=10, - protocols=["TCP", "UDP", "ICMP"], - ports=["HTTP", "DNS", "ARP"], - ip_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"], - act_map={}, - ) - observation_space = ObservationManager(NestedObservation(components={})) - reward_function = RewardFunction() - - test_agent = ControlledAgent( - agent_name="test_agent", - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - ) + test_agent = ControlledAgent(config=config) game.agents["test_agent"] = test_agent diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index fa4781db..4ca97a0e 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -49,7 +49,7 @@ def test_application_install_uninstall_on_uc2(): cfg = yaml.safe_load(f) env = PrimaiteGymEnv(env_config=cfg) - env.agent.flatten_obs = False + env.agent.config.agent_settings.flatten_obs = False env.reset() _, _, _, _, _ = env.step(0) diff --git a/tests/integration_tests/configuration_file_parsing/test_software_fix_duration.py b/tests/integration_tests/configuration_file_parsing/test_software_fixing_duration.py similarity index 73% rename from tests/integration_tests/configuration_file_parsing/test_software_fix_duration.py rename to tests/integration_tests/configuration_file_parsing/test_software_fixing_duration.py index b1c644cc..10896956 100644 --- a/tests/integration_tests/configuration_file_parsing/test_software_fix_duration.py +++ b/tests/integration_tests/configuration_file_parsing/test_software_fixing_duration.py @@ -13,8 +13,8 @@ from primaite.simulator.system.services.database.database_service import Databas from primaite.simulator.system.services.dns.dns_client import DNSClient from tests import TEST_ASSETS_ROOT -TEST_CONFIG = TEST_ASSETS_ROOT / "configs/software_fix_duration.yaml" -ONE_ITEM_CONFIG = TEST_ASSETS_ROOT / "configs/fix_duration_one_item.yaml" +TEST_CONFIG = TEST_ASSETS_ROOT / "configs/software_fixing_duration.yaml" +ONE_ITEM_CONFIG = TEST_ASSETS_ROOT / "configs/fixing_duration_one_item.yaml" TestApplications = ["DummyApplication", "BroadcastTestClient"] @@ -27,27 +27,27 @@ def load_config(config_path: Union[str, Path]) -> PrimaiteGame: return PrimaiteGame.from_config(cfg) -def test_default_fix_duration(): - """Test that software with no defined fix duration in config uses the default fix duration of 2.""" +def test_default_fixing_duration(): + """Test that software with no defined fixing duration in config uses the default fixing duration of 2.""" game = load_config(TEST_CONFIG) client_2: Computer = game.simulation.network.get_node_by_hostname("client_2") database_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient") - assert database_client.fixing_duration == 2 + assert database_client.config.fixing_duration == 2 dns_client: DNSClient = client_2.software_manager.software.get("DNSClient") - assert dns_client.fixing_duration == 2 + assert dns_client.config.fixing_duration == 2 -def test_fix_duration_set_from_config(): - """Test to check that the fix duration set for applications and services works as intended.""" +def test_fixing_duration_set_from_config(): + """Test to check that the fixing duration set for applications and services works as intended.""" game = load_config(TEST_CONFIG) client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") # in config - services take 3 timesteps to fix for service in ["DNSClient", "DNSServer", "DatabaseService", "WebServer", "FTPClient", "FTPServer", "NTPServer"]: assert client_1.software_manager.software.get(service) is not None - assert client_1.software_manager.software.get(service).fixing_duration == 3 + assert client_1.software_manager.software.get(service).config.fixing_duration == 3 # in config - applications take 1 timestep to fix # remove test applications from list @@ -55,27 +55,27 @@ def test_fix_duration_set_from_config(): for application in ["RansomwareScript", "WebBrowser", "DataManipulationBot", "DoSBot", "DatabaseClient"]: assert client_1.software_manager.software.get(application) is not None - assert client_1.software_manager.software.get(application).fixing_duration == 1 + assert client_1.software_manager.software.get(application).config.fixing_duration == 1 -def test_fix_duration_for_one_item(): - """Test that setting fix duration for one application does not affect other components.""" +def test_fixing_duration_for_one_item(): + """Test that setting fixing duration for one application does not affect other components.""" game = load_config(ONE_ITEM_CONFIG) client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") # in config - services take 3 timesteps to fix for service in ["DNSClient", "DNSServer", "WebServer", "FTPClient", "FTPServer", "NTPServer"]: assert client_1.software_manager.software.get(service) is not None - assert client_1.software_manager.software.get(service).fixing_duration == 2 + assert client_1.software_manager.software.get(service).config.fixing_duration == 2 # in config - applications take 1 timestep to fix # remove test applications from list for applications in ["RansomwareScript", "WebBrowser", "DataManipulationBot", "DoSBot"]: assert client_1.software_manager.software.get(applications) is not None - assert client_1.software_manager.software.get(applications).fixing_duration == 2 + assert client_1.software_manager.software.get(applications).config.fixing_duration == 2 database_client: DatabaseClient = client_1.software_manager.software.get("DatabaseClient") - assert database_client.fixing_duration == 1 + assert database_client.config.fixing_duration == 1 database_service: DatabaseService = client_1.software_manager.software.get("DatabaseService") - assert database_service.fixing_duration == 5 + assert database_service.config.fixing_duration == 5 diff --git a/tests/integration_tests/extensions/applications/extended_application.py b/tests/integration_tests/extensions/applications/extended_application.py index 9863dbba..159cfd06 100644 --- a/tests/integration_tests/extensions/applications/extended_application.py +++ b/tests/integration_tests/extensions/applications/extended_application.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address from typing import Dict, List, Optional from urllib.parse import urlparse -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from primaite import getLogger from primaite.interface.request import RequestResponse @@ -31,6 +31,14 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"): The application requests and loads web pages using its domain name and requesting IP addresses using DNS. """ + class ConfigSchema(Application.ConfigSchema): + """ConfigSchema for ExtendedApplication.""" + + type: str = "ExtendedApplication" + target_url: Optional[str] = None + + config: "ExtendedApplication.ConfigSchema" = Field(default_factory=lambda: ExtendedApplication.ConfigSchema()) + target_url: Optional[str] = None domain_name_ip_address: Optional[IPv4Address] = None @@ -50,6 +58,7 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"): kwargs["port"] = PORT_LOOKUP["HTTP"] super().__init__(**kwargs) + self.target_url = self.config.target_url self.run() def _init_request_manager(self) -> RequestManager: diff --git a/tests/integration_tests/extensions/services/extended_service.py b/tests/integration_tests/extensions/services/extended_service.py index 0924a91b..ba247369 100644 --- a/tests/integration_tests/extensions/services/extended_service.py +++ b/tests/integration_tests/extensions/services/extended_service.py @@ -3,6 +3,8 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Literal, Optional, Union from uuid import uuid4 +from pydantic import Field + from primaite import getLogger from primaite.simulator.file_system.file_system import File from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus @@ -17,13 +19,20 @@ from primaite.utils.validation.port import PORT_LOOKUP _LOGGER = getLogger(__name__) -class ExtendedService(Service, identifier="extendedservice"): +class ExtendedService(Service, identifier="ExtendedService"): """ A copy of DatabaseService that uses the extension framework instead of being part of PrimAITE. This class inherits from the `Service` class and provides methods to simulate a SQL database. """ + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for ExtendedService.""" + + type: str = "ExtendedService" + + config: "ExtendedService.ConfigSchema" = Field(default_factory=lambda: ExtendedService.ConfigSchema()) + password: Optional[str] = None """Password that needs to be provided by clients if they want to connect to the DatabaseService.""" diff --git a/tests/integration_tests/game_layer/actions/test_application_request_permission.py b/tests/integration_tests/game_layer/actions/test_application_request_permission.py index e90fa591..c47b617b 100644 --- a/tests/integration_tests/game_layer/actions/test_application_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_application_request_permission.py @@ -33,22 +33,22 @@ def test_application_cannot_perform_actions_unless_running(game_and_agent_fixtur browser.close() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0}) + action = ("node_application_scan", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("NODE_APPLICATION_CLOSE", {"node_id": 0, "application_id": 0}) + action = ("node_application_close", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("NODE_APPLICATION_FIX", {"node_id": 0, "application_id": 0}) + action = ("node_application_fix", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}) + action = ("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED diff --git a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py index 7930304e..59eb8a60 100644 --- a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py +++ b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py @@ -46,23 +46,21 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen server_1: Server = game.simulation.network.get_node_by_hostname("server_1") action = ( - "NODE_APPLICATION_INSTALL", - {"node_id": 1, "application_name": "C2Beacon"}, + "node_application_install", + {"node_name": "server_1", "application_name": "C2Beacon"}, ) agent.store_action(action) game.step() assert agent.history[-1].response.status == "success" action = ( - "CONFIGURE_C2_BEACON", + "configure_c2_beacon", { - "node_id": 1, - "config": { - "c2_server_ip_address": "10.0.1.2", - "keep_alive_frequency": 5, - "masquerade_protocol": "TCP", - "masquerade_port": "HTTP", - }, + "node_name": "server_1", + "c2_server_ip_address": "10.0.1.2", + "keep_alive_frequency": 5, + "masquerade_protocol": "TCP", + "masquerade_port": "HTTP", }, ) agent.store_action(action) @@ -70,8 +68,8 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen assert agent.history[-1].response.status == "success" action = ( - "NODE_APPLICATION_EXECUTE", - {"node_id": 1, "application_id": 0}, + "node_application_execute", + {"node_name": "server_1", "application_name": "C2Beacon"}, ) agent.store_action(action) game.step() @@ -103,14 +101,12 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA # C2 Action 1: Installing the RansomwareScript & Database client via Terminal action = ( - "C2_SERVER_TERMINAL_COMMAND", + "c2_server_terminal_command", { - "node_id": 0, + "node_name": "client_1", "ip_address": None, - "account": { - "username": "admin", - "password": "admin", - }, + "username": "admin", + "password": "admin", "commands": [ ["software_manager", "application", "install", "RansomwareScript"], ["software_manager", "application", "install", "DatabaseClient"], @@ -122,10 +118,11 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA assert agent.history[-1].response.status == "success" action = ( - "C2_SERVER_RANSOMWARE_CONFIGURE", + "c2_server_ransomware_configure", { - "node_id": 0, - "config": {"server_ip_address": "10.0.2.3", "payload": "ENCRYPT"}, + "node_name": "client_1", + "server_ip_address": "10.0.2.3", + "payload": "ENCRYPT", }, ) agent.store_action(action) @@ -134,16 +131,16 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA # Stepping a few timesteps to allow for the RansowmareScript to finish installing. - action = ("DONOTHING", {}) + action = ("do_nothing", {}) agent.store_action(action) game.step() game.step() game.step() action = ( - "C2_SERVER_RANSOMWARE_LAUNCH", + "c2_server_ransomware_launch", { - "node_id": 0, + "node_name": "client_1", }, ) agent.store_action(action) @@ -181,17 +178,15 @@ def test_c2_server_data_exfiltration(game_and_agent_fixture: Tuple[PrimaiteGame, # C2 Action: Data exfiltrate. action = ( - "C2_SERVER_DATA_EXFILTRATE", + "c2_server_data_exfiltrate", { - "node_id": 0, + "node_name": "client_1", "target_file_name": "database.db", "target_folder_name": "database", "exfiltration_folder_name": "spoils", "target_ip_address": "10.0.2.3", - "account": { - "username": "admin", - "password": "admin", - }, + "username": "admin", + "password": "admin", }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_configure_actions.py b/tests/integration_tests/game_layer/actions/test_configure_actions.py index 4e72e60d..17559405 100644 --- a/tests/integration_tests/game_layer/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address import pytest from pydantic import ValidationError -from primaite.game.agent.actions import ( +from primaite.game.agent.actions.software import ( ConfigureDatabaseClientAction, ConfigureDoSBotAction, ConfigureRansomwareScriptAction, @@ -27,7 +27,6 @@ class TestConfigureDatabaseAction: def test_configure_ip_password(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -35,13 +34,11 @@ class TestConfigureDatabaseAction: db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] action = ( - "CONFIGURE_DATABASE_CLIENT", + "configure_database_client", { - "node_id": 0, - "config": { - "server_ip_address": "192.168.1.99", - "server_password": "admin123", - }, + "node_name": "client_1", + "server_ip_address": "192.168.1.99", + "server_password": "admin123", }, ) agent.store_action(action) @@ -53,7 +50,6 @@ class TestConfigureDatabaseAction: def test_configure_ip(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -61,12 +57,10 @@ class TestConfigureDatabaseAction: db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] action = ( - "CONFIGURE_DATABASE_CLIENT", + "configure_database_client", { - "node_id": 0, - "config": { - "server_ip_address": "192.168.1.99", - }, + "node_name": "client_1", + "server_ip_address": "192.168.1.99", }, ) agent.store_action(action) @@ -78,7 +72,6 @@ class TestConfigureDatabaseAction: def test_configure_password(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -87,12 +80,10 @@ class TestConfigureDatabaseAction: old_ip = db_client.server_ip_address action = ( - "CONFIGURE_DATABASE_CLIENT", + "configure_database_client", { - "node_id": 0, - "config": { - "server_password": "admin123", - }, + "node_name": "client_1", + "server_password": "admin123", }, ) agent.store_action(action) @@ -120,9 +111,6 @@ class TestConfigureRansomwareScriptAction: def test_configure_ip_password(self, game_and_agent, config): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( - agent.action_manager - ) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -134,8 +122,8 @@ class TestConfigureRansomwareScriptAction: old_payload = ransomware_script.payload action = ( - "CONFIGURE_RANSOMWARE_SCRIPT", - {"node_id": 0, "config": config}, + "configure_ransomware_script", + {"node_name": "client_1", **config}, ) agent.store_action(action) game.step() @@ -151,18 +139,15 @@ class TestConfigureRansomwareScriptAction: def test_invalid_config(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction( - agent.action_manager - ) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") client_1.software_manager.install(RansomwareScript) ransomware_script: RansomwareScript = client_1.software_manager.software["RansomwareScript"] action = ( - "CONFIGURE_RANSOMWARE_SCRIPT", + "configure_ransomware_script", { - "node_id": 0, + "node_name": "client_1", "config": {"server_password": "admin123", "bad_option": 70}, }, ) @@ -172,28 +157,25 @@ class TestConfigureRansomwareScriptAction: class TestConfigureDoSBot: - def test_configure_DoSBot(self, game_and_agent): + def test_configure_dos_bot(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["CONFIGURE_DOSBOT"] = ConfigureDoSBotAction(agent.action_manager) client_1 = game.simulation.network.get_node_by_hostname("client_1") client_1.software_manager.install(DoSBot) dos_bot: DoSBot = client_1.software_manager.software["DoSBot"] action = ( - "CONFIGURE_DOSBOT", + "configure_dos_bot", { - "node_id": 0, - "config": { - "target_ip_address": "192.168.1.99", - "target_port": "POSTGRES_SERVER", - "payload": "HACC", - "repeat": False, - "port_scan_p_of_success": 0.875, - "dos_intensity": 0.75, - "max_sessions": 50, - }, + "node_name": "client_1", + "target_ip_address": "192.168.1.99", + "target_port": "POSTGRES_SERVER", + "payload": "HACC", + "repeat": False, + "port_scan_p_of_success": 0.875, + "dos_intensity": 0.75, + "max_sessions": 50, }, ) agent.store_action(action) @@ -239,7 +221,7 @@ class TestConfigureYAML: assert db_client.server_password == "correct_password" assert db_client.connect() - def test_configure_ransomware_script(self): + def test_c2_server_ransomware_configure(self): env = PrimaiteGymEnv(env_config=APP_CONFIG_YAML) client_2 = env.game.simulation.network.get_node_by_hostname("client_2") assert client_2.software_manager.software.get("RansomwareScript") is None diff --git a/tests/integration_tests/game_layer/actions/test_file_request_permission.py b/tests/integration_tests/game_layer/actions/test_file_request_permission.py index 91aa9fcd..cab80434 100644 --- a/tests/integration_tests/game_layer/actions/test_file_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_file_request_permission.py @@ -33,8 +33,8 @@ def test_create_file(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): assert client_1.file_system.get_file(folder_name=random_folder, file_name=random_file) is None action = ( - "NODE_FILE_CREATE", - {"node_id": 0, "folder_name": random_folder, "file_name": random_file}, + "node_file_create", + {"node_name": "client_1", "folder_name": random_folder, "file_name": random_file}, ) agent.store_action(action) game.step() @@ -51,8 +51,8 @@ def test_file_delete_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge assert file.deleted is False action = ( - "NODE_FILE_DELETE", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + "node_file_delete", + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -69,11 +69,11 @@ def test_file_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent file.corrupt() assert file.health_status == FileSystemItemHealthStatus.CORRUPT - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE action = ( - "NODE_FILE_SCAN", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + "node_file_scan", + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -93,8 +93,8 @@ def test_file_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge assert file.health_status == FileSystemItemHealthStatus.CORRUPT action = ( - "NODE_FILE_REPAIR", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + "node_file_repair", + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -113,8 +113,8 @@ def test_file_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg assert file.health_status == FileSystemItemHealthStatus.CORRUPT action = ( - "NODE_FILE_RESTORE", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + "node_file_restore", + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -132,8 +132,8 @@ def test_file_corrupt_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg assert file.health_status == FileSystemItemHealthStatus.GOOD action = ( - "NODE_FILE_CORRUPT", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + "node_file_corrupt", + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -150,8 +150,8 @@ def test_file_access_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge assert file.num_access == 0 action = ( - "NODE_FILE_ACCESS", - {"node_id": 0, "folder_name": file.folder_name, "file_name": file.name}, + "node_file_access", + {"node_name": "client_1", "folder_name": file.folder_name, "file_name": file.name}, ) agent.store_action(action) game.step() diff --git a/tests/integration_tests/game_layer/actions/test_folder_request_permission.py b/tests/integration_tests/game_layer/actions/test_folder_request_permission.py index 56bbbd4e..207f7d48 100644 --- a/tests/integration_tests/game_layer/actions/test_folder_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_folder_request_permission.py @@ -32,9 +32,9 @@ def test_create_folder(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): assert client_1.file_system.get_folder(folder_name=random_folder) is None action = ( - "NODE_FOLDER_CREATE", + "node_folder_create", { - "node_id": 0, + "node_name": "client_1", "folder_name": random_folder, }, ) @@ -52,18 +52,18 @@ def test_folder_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge folder = client_1.file_system.get_folder(folder_name="downloads") assert folder.health_status == FileSystemItemHealthStatus.GOOD - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE folder.corrupt() assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE action = ( - "NODE_FOLDER_SCAN", + "node_folder_scan", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) @@ -87,10 +87,10 @@ def test_folder_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA assert folder.health_status == FileSystemItemHealthStatus.CORRUPT action = ( - "NODE_FOLDER_REPAIR", + "node_folder_repair", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) @@ -111,10 +111,10 @@ def test_folder_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy assert folder.health_status == FileSystemItemHealthStatus.CORRUPT action = ( - "NODE_FOLDER_RESTORE", + "node_folder_restore", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_nic_request_permission.py b/tests/integration_tests/game_layer/actions/test_nic_request_permission.py index 8846809d..11e39c7e 100644 --- a/tests/integration_tests/game_layer/actions/test_nic_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_nic_request_permission.py @@ -29,10 +29,10 @@ def test_nic_cannot_be_turned_off_if_not_on(game_and_agent_fixture: Tuple[Primai assert nic.enabled is False action = ( - "HOST_NIC_DISABLE", + "host_nic_disable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -50,10 +50,10 @@ def test_nic_cannot_be_turned_on_if_already_on(game_and_agent_fixture: Tuple[Pri assert nic.enabled action = ( - "HOST_NIC_ENABLE", + "host_nic_enable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -71,10 +71,10 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr assert nic.enabled action = ( - "HOST_NIC_DISABLE", + "host_nic_disable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -83,10 +83,10 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr assert nic.enabled is False action = ( - "HOST_NIC_ENABLE", + "host_nic_enable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_node_request_permission.py b/tests/integration_tests/game_layer/actions/test_node_request_permission.py index d75567e9..8a438673 100644 --- a/tests/integration_tests/game_layer/actions/test_node_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_node_request_permission.py @@ -29,28 +29,28 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy assert client_1.operating_state == NodeOperatingState.ON # turn it off - action = ("NODE_SHUTDOWN", {"node_id": 0}) + action = ("node_shutdown", {"node_name": "client_1"}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN for i in range(client_1.shut_down_duration + 1): - action = ("DONOTHING", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.OFF # turn it on - action = ("NODE_STARTUP", {"node_id": 0}) + action = ("node_startup", {"node_name": "client_1"}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.BOOTING for i in range(client_1.start_up_duration + 1): - action = ("DONOTHING", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() @@ -65,7 +65,7 @@ def test_node_cannot_be_started_up_if_node_is_already_on(game_and_agent_fixture: assert client_1.operating_state == NodeOperatingState.ON # turn it on - action = ("NODE_STARTUP", {"node_id": 0}) + action = ("node_startup", {"node_name": "client_1"}) agent.store_action(action) game.step() @@ -80,14 +80,14 @@ def test_node_cannot_be_shut_down_if_node_is_already_off(game_and_agent_fixture: client_1.power_off() for i in range(client_1.shut_down_duration + 1): - action = ("DONOTHING", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.OFF # turn it ff - action = ("NODE_SHUTDOWN", {"node_id": 0}) + action = ("node_shutdown", {"node_name": "client_1"}) agent.store_action(action) game.step() diff --git a/tests/integration_tests/game_layer/actions/test_service_request_permission.py b/tests/integration_tests/game_layer/actions/test_service_request_permission.py index ebc9fd3b..80e68131 100644 --- a/tests/integration_tests/game_layer/actions/test_service_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_service_request_permission.py @@ -31,7 +31,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): dns_server.pause() assert dns_server.operating_state == ServiceOperatingState.PAUSED - action = ("NODE_SERVICE_START", {"node_id": 1, "service_id": 0}) + action = ("node_service_start", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.PAUSED @@ -40,7 +40,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_START", {"node_id": 1, "service_id": 0}) + action = ("node_service_start", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -54,7 +54,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]) server_1: Server = game.simulation.network.get_node_by_hostname("server_1") dns_server = server_1.software_manager.software.get("DNSServer") - action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.RUNNING @@ -63,7 +63,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]) assert dns_server.operating_state == ServiceOperatingState.PAUSED - action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -80,27 +80,27 @@ def test_service_cannot_perform_actions_unless_running(game_and_agent_fixture: T dns_server.stop() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_PAUSE", {"node_id": 1, "service_id": 0}) + action = ("node_service_pause", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_RESTART", {"node_id": 1, "service_id": 0}) + action = ("node_service_restart", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("NODE_SERVICE_FIX", {"node_id": 1, "service_id": 0}) + action = ("node_service_fix", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED diff --git a/tests/integration_tests/game_layer/actions/test_terminal_actions.py b/tests/integration_tests/game_layer/actions/test_terminal_actions.py index 96110656..f15f7156 100644 --- a/tests/integration_tests/game_layer/actions/test_terminal_actions.py +++ b/tests/integration_tests/game_layer/actions/test_terminal_actions.py @@ -36,9 +36,9 @@ def test_remote_login(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): server_1_usm.add_user("user123", "password", is_admin=True) action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -68,9 +68,9 @@ def test_remote_login_wrong_password(game_and_agent_fixture: Tuple[PrimaiteGame, server_1_usm.add_user("user123", "password", is_admin=True) action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "wrong_password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -100,12 +100,13 @@ def test_remote_login_change_password(game_and_agent_fixture: Tuple[PrimaiteGame server_1_um.add_user("user123", "password", is_admin=True) action = ( - "NODE_ACCOUNTS_CHANGE_PASSWORD", + "node_account_change_password", { - "node_id": 1, # server_1 + "node_name": "server_1", # server_1 "username": "user123", "current_password": "password", "new_password": "different_password", + "remote_ip": str(server_1.network_interface[1].ip_address), }, ) agent.store_action(action) @@ -126,9 +127,9 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam # Log in remotely action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -139,12 +140,13 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam # Change password action = ( - "NODE_ACCOUNTS_CHANGE_PASSWORD", + "node_account_change_password", { - "node_id": 1, # server_1 + "node_name": "server_1", # server_1 "username": "user123", "current_password": "password", "new_password": "different_password", + "remote_ip": str(server_1.network_interface[1].ip_address), }, ) agent.store_action(action) @@ -152,9 +154,9 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam # Assert that the user cannot execute an action action = ( - "NODE_SEND_REMOTE_COMMAND", + "node_send_remote_command", { - "node_id": 0, + "node_name": "client_1", "remote_ip": str(server_1.network_interface[1].ip_address), "command": ["file_system", "create", "file", "folder123", "doggo.pdf", False], }, diff --git a/tests/integration_tests/game_layer/observations/test_file_system_observations.py b/tests/integration_tests/game_layer/observations/test_file_system_observations.py index a56deb2b..7323461c 100644 --- a/tests/integration_tests/game_layer/observations/test_file_system_observations.py +++ b/tests/integration_tests/game_layer/observations/test_file_system_observations.py @@ -32,11 +32,11 @@ def test_file_observation(simulation): assert dog_file_obs.space["health_status"] == spaces.Discrete(6) observation_state = dog_file_obs.observe(simulation.describe_state()) - assert observation_state.get("health_status") == 1 # good initial + assert observation_state.get("health_status") == 0 # initially unset file.corrupt() observation_state = dog_file_obs.observe(simulation.describe_state()) - assert observation_state.get("health_status") == 1 # scan file so this changes + assert observation_state.get("health_status") == 0 # still default unset value because no scan happened file.scan() file.apply_timestep(0) # apply time step @@ -63,11 +63,11 @@ def test_folder_observation(simulation): observation_state = root_folder_obs.observe(simulation.describe_state()) assert observation_state.get("FILES") is not None - assert observation_state.get("health_status") == 1 + assert observation_state.get("health_status") == 0 # initially unset file.corrupt() # corrupt just the file observation_state = root_folder_obs.observe(simulation.describe_state()) - assert observation_state.get("health_status") == 1 # scan folder to change this + assert observation_state.get("health_status") == 0 # still unset as no scan occurred yet folder.scan() for i in range(folder.scan_duration + 1): diff --git a/tests/integration_tests/game_layer/observations/test_nic_observations.py b/tests/integration_tests/game_layer/observations/test_nic_observations.py index 0ad03198..bd9417ba 100644 --- a/tests/integration_tests/game_layer/observations/test_nic_observations.py +++ b/tests/integration_tests/game_layer/observations/test_nic_observations.py @@ -191,7 +191,7 @@ def test_nic_monitored_traffic(simulation): # send a database query browser: WebBrowser = pc.software_manager.software.get("WebBrowser") - browser.target_url = f"http://arcd.com/" + browser.config.target_url = f"http://arcd.com/" browser.get_webpage() traffic_obs = nic_obs.observe(simulation.describe_state()).get("TRAFFIC") diff --git a/tests/integration_tests/game_layer/observations/test_user_observations.py b/tests/integration_tests/game_layer/observations/test_user_observations.py index 92c533c9..de5c3c98 100644 --- a/tests/integration_tests/game_layer/observations/test_user_observations.py +++ b/tests/integration_tests/game_layer/observations/test_user_observations.py @@ -13,7 +13,7 @@ DATA_MANIPULATION_CONFIG = TEST_ASSETS_ROOT / "configs" / "data_manipulation.yam def env_with_ssh() -> PrimaiteGymEnv: """Build data manipulation environment with SSH port open on router.""" env = PrimaiteGymEnv(DATA_MANIPULATION_CONFIG) - env.agent.flatten_obs = False + env.agent.config.agent_settings.flatten_obs = False router: Router = env.game.simulation.network.get_node_by_hostname("router_1") router.acl.add_rule(ACLAction.PERMIT, src_port=PORT_LOOKUP["SSH"], dst_port=PORT_LOOKUP["SSH"], position=3) return env diff --git a/tests/integration_tests/game_layer/test_RNG_seed.py b/tests/integration_tests/game_layer/test_RNG_seed.py index 45ad2708..464f95db 100644 --- a/tests/integration_tests/game_layer/test_RNG_seed.py +++ b/tests/integration_tests/game_layer/test_RNG_seed.py @@ -24,12 +24,12 @@ def test_rng_seed_set(create_env): env.reset(seed=3) for i in range(100): env.step(0) - a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"] + a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"] env.reset(seed=3) for i in range(100): env.step(0) - b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"] + b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"] assert a == b @@ -40,11 +40,11 @@ def test_rng_seed_unset(create_env): env.reset() for i in range(100): env.step(0) - a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"] + a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"] env.reset() for i in range(100): env.step(0) - b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"] + b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"] assert a != b diff --git a/tests/integration_tests/game_layer/test_action_mask.py b/tests/integration_tests/game_layer/test_action_mask.py index 3d26b73d..75965f16 100644 --- a/tests/integration_tests/game_layer/test_action_mask.py +++ b/tests/integration_tests/game_layer/test_action_mask.py @@ -15,7 +15,6 @@ def test_mask_contents_correct(): net = sim.network mask = game.action_mask("defender") agent = env.agent - node_list = agent.action_manager.node_names action_map = agent.action_manager.action_map # CHECK NIC ENABLE/DISABLE ACTIONS @@ -23,10 +22,10 @@ def test_mask_contents_correct(): mask = game.action_mask("defender") act_type, act_params = action - if act_type == "NODE_NIC_ENABLE": - node_name = node_list[act_params["node_id"]] + if act_type == "node_nic_enable": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) - nic_obj = node_obj.network_interface[act_params["nic_id"] + 1] + nic_obj = node_obj.network_interface[act_params["nic_num"]] assert nic_obj.enabled assert not mask[action_num] nic_obj.disable() @@ -34,10 +33,10 @@ def test_mask_contents_correct(): assert mask[action_num] nic_obj.enable() - if act_type == "NODE_NIC_DISABLE": - node_name = node_list[act_params["node_id"]] + if act_type == "node_nic_disable": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) - nic_obj = node_obj.network_interface[act_params["nic_id"] + 1] + nic_obj = node_obj.network_interface[act_params["nic_num"]] assert nic_obj.enabled assert mask[action_num] nic_obj.disable() @@ -45,14 +44,14 @@ def test_mask_contents_correct(): assert not mask[action_num] nic_obj.enable() - if act_type == "ROUTER_ACL_ADDRULE": + if act_type == "router_acl_add_rule": assert mask[action_num] - if act_type == "ROUTER_ACL_REMOVERULE": + if act_type == "router_acl_remove_rule": assert mask[action_num] - if act_type == "NODE_RESET": - node_name = node_list[act_params["node_id"]] + if act_type == "node_reset": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) assert node_obj.operating_state is NodeOperatingState.ON assert mask[action_num] @@ -61,8 +60,8 @@ def test_mask_contents_correct(): assert not mask[action_num] node_obj.operating_state = NodeOperatingState.ON - if act_type == "NODE_SHUTDOWN": - node_name = node_list[act_params["node_id"]] + if act_type == "node_shutdown": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) assert node_obj.operating_state is NodeOperatingState.ON assert mask[action_num] @@ -71,8 +70,8 @@ def test_mask_contents_correct(): assert not mask[action_num] node_obj.operating_state = NodeOperatingState.ON - if act_type == "NODE_OS_SCAN": - node_name = node_list[act_params["node_id"]] + if act_type == "node_os_scan": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) assert node_obj.operating_state is NodeOperatingState.ON assert mask[action_num] @@ -81,8 +80,8 @@ def test_mask_contents_correct(): assert not mask[action_num] node_obj.operating_state = NodeOperatingState.ON - if act_type == "NODE_STARTUP": - node_name = node_list[act_params["node_id"]] + if act_type == "node_startup": + node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) assert node_obj.operating_state is NodeOperatingState.ON assert not mask[action_num] @@ -91,15 +90,15 @@ def test_mask_contents_correct(): assert mask[action_num] node_obj.operating_state = NodeOperatingState.ON - if act_type == "DONOTHING": + if act_type == "do_nothing": assert mask[action_num] - if act_type == "NODE_SERVICE_DISABLE": + if act_type == "node_service_disable": assert mask[action_num] - if act_type in ["NODE_SERVICE_SCAN", "NODE_SERVICE_STOP", "NODE_SERVICE_PAUSE"]: - node_name = node_list[act_params["node_id"]] - service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]] + if act_type in ["node_service_scan", "node_service_stop", "node_service_pause"]: + node_name = act_params["node_name"] + service_name = act_params["service_name"] node_obj = net.get_node_by_hostname(node_name) service_obj = node_obj.software_manager.software.get(service_name) assert service_obj.operating_state is ServiceOperatingState.RUNNING @@ -109,9 +108,9 @@ def test_mask_contents_correct(): assert not mask[action_num] service_obj.operating_state = ServiceOperatingState.RUNNING - if act_type == "NODE_SERVICE_RESUME": - node_name = node_list[act_params["node_id"]] - service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]] + if act_type == "node_service_resume": + node_name = act_params["node_name"] + service_name = act_params["service_name"] node_obj = net.get_node_by_hostname(node_name) service_obj = node_obj.software_manager.software.get(service_name) assert service_obj.operating_state is ServiceOperatingState.RUNNING @@ -121,9 +120,9 @@ def test_mask_contents_correct(): assert mask[action_num] service_obj.operating_state = ServiceOperatingState.RUNNING - if act_type == "NODE_SERVICE_START": - node_name = node_list[act_params["node_id"]] - service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]] + if act_type == "node_service_start": + node_name = act_params["node_name"] + service_name = act_params["service_name"] node_obj = net.get_node_by_hostname(node_name) service_obj = node_obj.software_manager.software.get(service_name) assert service_obj.operating_state is ServiceOperatingState.RUNNING @@ -133,9 +132,9 @@ def test_mask_contents_correct(): assert mask[action_num] service_obj.operating_state = ServiceOperatingState.RUNNING - if act_type == "NODE_SERVICE_ENABLE": - node_name = node_list[act_params["node_id"]] - service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]] + if act_type == "node_service_enable": + node_name = act_params["node_name"] + service_name = act_params["service_name"] node_obj = net.get_node_by_hostname(node_name) service_obj = node_obj.software_manager.software.get(service_name) assert service_obj.operating_state is ServiceOperatingState.RUNNING @@ -145,12 +144,10 @@ def test_mask_contents_correct(): assert mask[action_num] service_obj.operating_state = ServiceOperatingState.RUNNING - if act_type in ["NODE_FILE_SCAN", "NODE_FILE_CHECKHASH", "NODE_FILE_DELETE"]: - node_name = node_list[act_params["node_id"]] - folder_name = agent.action_manager.get_folder_name_by_idx(act_params["node_id"], act_params["folder_id"]) - file_name = agent.action_manager.get_file_name_by_idx( - act_params["node_id"], act_params["folder_id"], act_params["file_id"] - ) + if act_type in ["node_file_scan", "node_file_checkhash", "node_file_delete"]: + node_name = act_params["node_name"] + folder_name = act_params["folder_name"] + file_name = act_params["file_name"] node_obj = net.get_node_by_hostname(node_name) file_obj = node_obj.file_system.get_file(folder_name, file_name, include_deleted=True) assert not file_obj.deleted diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index beb7b6a8..5a308cf8 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -32,10 +32,10 @@ FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network. def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): - """Test that the DoNothingAction can form a request and that it is accepted by the simulation.""" + """Test that the do_nothingAction can form a request and that it is accepted by the simulation.""" game, agent = game_and_agent - action = ("DONOTHING", {}) + action = ("do_nothing", {}) agent.store_action(action) game.step() @@ -56,7 +56,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.UNUSED # 2: Scan and check that the visible state is now correct - action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.GOOD @@ -67,7 +67,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.GOOD # 4: Scan and check that the visible state is now correct - action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.COMPROMISED @@ -88,7 +88,7 @@ def test_node_service_fix_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA svc.health_state_actual = SoftwareHealthState.COMPROMISED # 2: Apply a patch action - action = ("NODE_SERVICE_FIX", {"node_id": 1, "service_id": 0}) + action = ("node_service_fix", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -96,7 +96,7 @@ def test_node_service_fix_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA assert svc.health_state_actual == SoftwareHealthState.FIXING # 4: perform a few do-nothing steps and check that the service is now in the good state - action = ("DONOTHING", {}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.GOOD @@ -121,18 +121,18 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox # 2: Add a rule to block client 1 from reaching server 2 on router action = ( - "ROUTER_ACL_ADDRULE", + "router_acl_add_rule", { "target_router": "router", - "position": 4, # 4th rule - "permission": 2, # DENY - "source_ip_id": 3, # 10.0.1.2 (client_1) - "dest_ip_id": 6, # 10.0.2.3 (server_2) - "dest_port_id": 1, # ALL - "source_port_id": 1, # ALL - "protocol_id": 1, # ALL - "source_wildcard_id": 0, - "dest_wildcard_id": 0, + "position": 4, + "permission": "DENY", + "src_ip": "10.0.1.2", + "src_wildcard": "NONE", + "src_port": "ALL", + "dst_ip": "10.0.2.3", + "dst_wildcard": "NONE", + "dst_port": "ALL", + "protocol_name": "icmp", }, ) agent.store_action(action) @@ -148,24 +148,26 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox # 4: Add a rule to block server_1 from reaching server_2 on router (this should not affect comms as they are on same subnet) action = ( - "ROUTER_ACL_ADDRULE", + "router_acl_add_rule", { "target_router": "router", "position": 5, # 5th rule - "permission": 2, # DENY - "source_ip_id": 5, # 10.0.2.2 (server_1) - "dest_ip_id": 6, # 10.0.2.3 (server_2) - "dest_port_id": 1, # ALL - "source_port_id": 1, # ALL - "protocol_id": 1, # ALL - "source_wildcard_id": 0, - "dest_wildcard_id": 0, + "permission": "DENY", # DENY + "src_ip": "10.0.2.2", # 10.0.2.2 (server_1) + "src_wildcard": 0, + "src_port": "ALL", # ALL + "dst_ip": "10.0.2.3", # 10.0.2.3 (server_2) + "dst_wildcard": 0, + "dst_port": "ALL", # ALL + "protocol_name": "ALL", # ALL }, ) agent.store_action(action) + print(agent.most_recent_action) game.step() - + print(agent.most_recent_action) # 5: Check that the ACL now has 6 rules, but that server_1 can still ping server_2 + print(router.acl.show()) assert router.acl.num_rules == 6 assert server_1.ping("10.0.2.3") # Can ping server_2 @@ -181,12 +183,12 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") browser.run() - browser.target_url = "http://www.example.com" + browser.config.target_url = "http://www.example.com" assert browser.get_webpage() # check that the browser can access example.com before we block it # 2: Remove rule that allows HTTP traffic across the network action = ( - "ROUTER_ACL_REMOVERULE", + "router_acl_remove_rule", { "target_router": "router", "position": 3, # 4th rule @@ -214,15 +216,15 @@ def test_host_nic_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") browser.run() - browser.target_url = "http://www.example.com" + browser.config.target_url = "http://www.example.com" assert browser.get_webpage() # check that the browser can access example.com before we block it # 2: Disable the NIC on client_1 action = ( - "HOST_NIC_DISABLE", + "host_nic_disable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -250,10 +252,10 @@ def test_host_nic_enable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAg # 2: Use action to enable nic action = ( - "HOST_NIC_ENABLE", + "host_nic_enable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -273,15 +275,15 @@ def test_node_file_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAge client_1 = game.simulation.network.get_node_by_hostname("client_1") file = client_1.file_system.get_file("downloads", "cat.png") assert file.health_status == FileSystemItemHealthStatus.GOOD - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE # 2: perform a scan and make sure nothing has changed action = ( - "NODE_FILE_SCAN", + "node_file_scan", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads, - "file_id": 0, # cat.png + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads, + "file_name": "cat.png", # cat.png }, ) agent.store_action(action) @@ -314,11 +316,11 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA # 2: delete the file action = ( - "NODE_FILE_DELETE", + "node_file_delete", { - "node_id": 0, # client_1 - "folder_id": 0, # downloads - "file_id": 0, # cat.png + "node_name": "client_1", # client_1 + "folder_name": "downloads", # downloads + "file_name": "cat.png", # cat.png }, ) agent.store_action(action) @@ -334,14 +336,15 @@ def test_node_file_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): """Test that a file is created.""" game, agent = game_and_agent - client_1 = game.simulation.network.get_node_by_hostname("client_1") # + client_1 = game.simulation.network.get_node_by_hostname("client_1") action = ( - "NODE_FILE_CREATE", + "node_file_create", { - "node_id": 0, + "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", + "force": "False", }, ) agent.store_action(action) @@ -357,9 +360,9 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): client_1 = game.simulation.network.get_node_by_hostname("client_1") # action = ( - "NODE_FILE_CREATE", + "node_file_create", { - "node_id": 0, + "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", }, @@ -370,9 +373,9 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 0 action = ( - "NODE_FILE_ACCESS", + "node_file_access", { - "node_id": 0, + "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", }, @@ -390,9 +393,9 @@ def test_node_folder_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): client_1 = game.simulation.network.get_node_by_hostname("client_1") # action = ( - "NODE_FOLDER_CREATE", + "node_folder_create", { - "node_id": 0, + "node_name": "client_1", "folder_name": "test", }, ) @@ -413,15 +416,15 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") browser.run() - browser.target_url = "http://www.example.com" + browser.config.target_url = "http://www.example.com" assert browser.get_webpage() # check that the browser can access example.com before we block it # 2: Disable the NIC on client_1 action = ( - "NETWORK_PORT_DISABLE", + "network_port_disable", { "target_nodename": "router", # router - "port_id": 1, # port 1 + "port_num": 1, # port 1 }, ) agent.store_action(action) @@ -450,10 +453,10 @@ def test_network_router_port_enable_integration(game_and_agent: Tuple[PrimaiteGa # 2: Use action to enable port action = ( - "NETWORK_PORT_ENABLE", + "network_port_enable", { "target_nodename": "router", # router - "port_id": 1, # port 1 + "port_num": 1, # port 1 }, ) agent.store_action(action) @@ -473,14 +476,17 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") browser.run() - browser.target_url = "http://www.example.com" + browser.config.target_url = "http://www.example.com" assert browser.get_webpage() # check that the browser can access example.com assert browser.health_state_actual == SoftwareHealthState.GOOD assert browser.health_state_visible == SoftwareHealthState.UNUSED # 2: Scan and check that the visible state is now correct - action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0}) + action = ( + "node_application_scan", + {"node_name": "client_1", "application_name": "WebBrowser"}, + ) agent.store_action(action) game.step() assert browser.health_state_actual == SoftwareHealthState.GOOD @@ -491,7 +497,10 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P assert browser.health_state_visible == SoftwareHealthState.GOOD # 4: Scan and check that the visible state is now correct - action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0}) + action = ( + "node_application_scan", + {"node_name": "client_1", "application_name": "WebBrowser"}, + ) agent.store_action(action) game.step() assert browser.health_state_actual == SoftwareHealthState.COMPROMISED @@ -512,7 +521,10 @@ def test_node_application_fix_integration(game_and_agent: Tuple[PrimaiteGame, Pr browser.health_state_actual = SoftwareHealthState.COMPROMISED # 2: Apply a fix action - action = ("NODE_APPLICATION_FIX", {"node_id": 0, "application_id": 0}) + action = ( + "node_application_fix", + {"node_name": "client_1", "application_name": "WebBrowser"}, + ) agent.store_action(action) game.step() @@ -520,7 +532,7 @@ def test_node_application_fix_integration(game_and_agent: Tuple[PrimaiteGame, Pr assert browser.health_state_actual == SoftwareHealthState.FIXING # 4: perform a few do-nothing steps and check that the application is now in the good state - action = ("DONOTHING", {}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert browser.health_state_actual == SoftwareHealthState.GOOD @@ -538,7 +550,10 @@ def test_node_application_close_integration(game_and_agent: Tuple[PrimaiteGame, assert browser.operating_state == ApplicationOperatingState.RUNNING # 2: Apply a close action - action = ("NODE_APPLICATION_CLOSE", {"node_id": 0, "application_id": 0}) + action = ( + "node_application_close", + {"node_name": "client_1", "application_name": "WebBrowser"}, + ) agent.store_action(action) game.step() @@ -549,7 +564,7 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl """Test that the NodeApplicationInstallAction and NodeApplicationRemoveAction can form a request and that it is accepted by the simulation. - When you initiate a install action, the Application will be installed and configured on the node. + When you initiate an install action, the Application will be installed and configured on the node. The remove action will uninstall the application from the node.""" game, agent = game_and_agent @@ -557,13 +572,19 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl assert client_1.software_manager.software.get("DoSBot") is None - action = ("NODE_APPLICATION_INSTALL", {"node_id": 0, "application_name": "DoSBot"}) + action = ( + "node_application_install", + {"node_name": "client_1", "application_name": "DoSBot"}, + ) agent.store_action(action) game.step() assert client_1.software_manager.software.get("DoSBot") is not None - action = ("NODE_APPLICATION_REMOVE", {"node_id": 0, "application_name": "DoSBot"}) + action = ( + "node_application_remove", + {"node_name": "client_1", "application_name": "DoSBot"}, + ) agent.store_action(action) game.step() @@ -656,9 +677,9 @@ def test_firewall_acl_add_remove_rule_integration(): assert firewall.external_outbound_acl.acl[1].action.name == "DENY" assert firewall.external_outbound_acl.acl[1].src_ip_address == IPv4Address("192.168.20.10") assert firewall.external_outbound_acl.acl[1].dst_ip_address == IPv4Address("192.168.0.10") - assert firewall.external_outbound_acl.acl[1].dst_port is None - assert firewall.external_outbound_acl.acl[1].src_port is None - assert firewall.external_outbound_acl.acl[1].protocol is None + assert firewall.external_outbound_acl.acl[1].dst_port == PORT_LOOKUP["NONE"] + assert firewall.external_outbound_acl.acl[1].src_port == PORT_LOOKUP["NONE"] + assert firewall.external_outbound_acl.acl[1].protocol == PROTOCOL_LOOKUP["NONE"] env.step(12) # Remove ACL rule from External Outbound assert firewall.external_outbound_acl.num_rules == 1 diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index 23364f13..5afad296 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -17,12 +17,7 @@ def test_file_observation(): dog_file_obs = FileObservation( where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"], include_num_access=False, - file_system_requires_scan=True, + file_system_requires_scan=False, ) assert dog_file_obs.observe(state) == {"health_status": 1} assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)}) - - -# TODO: -# def test_file_num_access(): -# ... diff --git a/tests/integration_tests/game_layer/test_rewards.py b/tests/integration_tests/game_layer/test_rewards.py index be55d915..de61b215 100644 --- a/tests/integration_tests/game_layer/test_rewards.py +++ b/tests/integration_tests/game_layer/test_rewards.py @@ -29,16 +29,16 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle client_1 = game.simulation.network.get_node_by_hostname("client_1") browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") browser.run() - browser.target_url = "http://www.example.com" + browser.config.target_url = "http://www.example.com" agent.reward_function.register_component(comp, 0.7) # Check that before trying to fetch the webpage, the reward is 0.0 - agent.store_action(("DONOTHING", {})) + agent.store_action(("do_nothing", {})) game.step() assert agent.reward_function.current_reward == 0.0 # Check that successfully fetching the webpage yields a reward of 0.7 - agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0})) + agent.store_action(("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"})) game.step() assert agent.reward_function.current_reward == 0.7 @@ -50,7 +50,7 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle src_port=PORT_LOOKUP["HTTP"], dst_port=PORT_LOOKUP["HTTP"], ) - agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0})) + agent.store_action(("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"})) game.step() assert agent.reward_function.current_reward == -0.7 @@ -83,7 +83,7 @@ def test_uc2_rewards(game_and_agent: tuple[PrimaiteGame, ControlledAgent]): response = game.simulation.apply_request(request) state = game.get_sim_state() ahi = AgentHistoryItem( - timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response + timestep=0, action="node_application_execute", parameters={}, request=request, response=response ) reward_value = comp.calculate(state, last_action_response=ahi) assert reward_value == 1.0 @@ -94,7 +94,7 @@ def test_uc2_rewards(game_and_agent: tuple[PrimaiteGame, ControlledAgent]): response = game.simulation.apply_request(request) state = game.get_sim_state() ahi = AgentHistoryItem( - timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response + timestep=0, action="node_application_execute", parameters={}, request=request, response=response ) reward_value = comp.calculate( state, @@ -154,13 +154,13 @@ def test_action_penalty(): # Penalty = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125) Penalty = ActionPenalty(config=schema) - # Assert that penalty is applied if action isn't DONOTHING + # Assert that penalty is applied if action isn't do_nothing reward_value = Penalty.calculate( state={}, last_action_response=AgentHistoryItem( timestep=0, - action="NODE_APPLICATION_EXECUTE", - parameters={"node_id": 0, "application_id": 1}, + action="node_application_execute", + parameters={"node_name": "client", "application_name": "WebBrowser"}, request=["execute"], response=RequestResponse.from_bool(True), ), @@ -168,12 +168,12 @@ def test_action_penalty(): assert reward_value == -0.75 - # Assert that no penalty applied for a DONOTHING action + # Assert that no penalty applied for a do_nothing action reward_value = Penalty.calculate( state={}, last_action_response=AgentHistoryItem( timestep=0, - action="DONOTHING", + action="do_nothing", parameters={}, request=["do_nothing"], response=RequestResponse.from_bool(True), @@ -192,12 +192,12 @@ def test_action_penalty_e2e(game_and_agent: tuple[PrimaiteGame, ControlledAgent] agent.reward_function.register_component(comp, 1.0) - action = ("DONOTHING", {}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert agent.reward_function.current_reward == 0.125 - action = ("NODE_FILE_SCAN", {"node_id": 0, "folder_id": 0, "file_id": 0}) + action = ("node_file_scan", {"node_name": "client", "folder_name": "downloads", "file_name": "document.pdf"}) agent.store_action(action) game.step() assert agent.reward_function.current_reward == -0.75 diff --git a/tests/integration_tests/network/test_broadcast.py b/tests/integration_tests/network/test_broadcast.py index 33fe70c3..ed40334f 100644 --- a/tests/integration_tests/network/test_broadcast.py +++ b/tests/integration_tests/network/test_broadcast.py @@ -3,6 +3,7 @@ from ipaddress import IPv4Address, IPv4Network from typing import Any, Dict, List, Tuple import pytest +from pydantic import Field from primaite.simulator.network.container import Network from primaite.simulator.network.hardware.nodes.host.computer import Computer @@ -14,9 +15,16 @@ from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP -class BroadcastTestService(Service): +class BroadcastTestService(Service, identifier="BroadcastTestService"): """A service for sending broadcast and unicast messages over a network.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for BroadcastTestService.""" + + type: str = "BroadcastTestService" + + config: "BroadcastTestService.ConfigSchema" = Field(default_factory=lambda: BroadcastTestService.ConfigSchema()) + def __init__(self, **kwargs): # Set default service properties for broadcasting kwargs["name"] = "BroadcastService" @@ -46,6 +54,13 @@ class BroadcastTestService(Service): class BroadcastTestClient(Application, identifier="BroadcastTestClient"): """A client application to receive broadcast and unicast messages.""" + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for BroadcastTestClient.""" + + type: str = "BroadcastTestClient" + + config: ConfigSchema = Field(default_factory=lambda: BroadcastTestClient.ConfigSchema()) + payloads_received: List = [] def __init__(self, **kwargs): diff --git a/tests/integration_tests/system/red_applications/test_c2_suite_integration.py b/tests/integration_tests/system/red_applications/test_c2_suite_integration.py index d88f8249..40226be6 100644 --- a/tests/integration_tests/system/red_applications/test_c2_suite_integration.py +++ b/tests/integration_tests/system/red_applications/test_c2_suite_integration.py @@ -495,6 +495,12 @@ def test_c2_suite_yaml(): computer_b: Computer = yaml_network.get_node_by_hostname("node_b") c2_beacon: C2Beacon = computer_b.software_manager.software.get("C2Beacon") + c2_beacon.configure( + c2_server_ip_address=c2_beacon.config.c2_server_ip_address, + keep_alive_frequency=c2_beacon.config.keep_alive_frequency, + masquerade_port=c2_beacon.config.masquerade_port, + masquerade_protocol=c2_beacon.config.masquerade_protocol, + ) assert c2_server.operating_state == ApplicationOperatingState.RUNNING diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index 674603fa..bb25f8c8 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -163,7 +163,7 @@ def test_restore_backup_without_updating_scan(uc2_network): db_service.db_file.corrupt() # corrupt the db assert db_service.db_file.health_status == FileSystemItemHealthStatus.CORRUPT # db file is actually corrupt - assert db_service.db_file.visible_health_status == FileSystemItemHealthStatus.GOOD # not scanned yet + assert db_service.db_file.visible_health_status == FileSystemItemHealthStatus.NONE # not scanned yet db_service.db_file.scan() # scan the db file @@ -190,7 +190,7 @@ def test_restore_backup_after_deleting_file_without_updating_scan(uc2_network): db_service.db_file.corrupt() # corrupt the db assert db_service.db_file.health_status == FileSystemItemHealthStatus.CORRUPT # db file is actually corrupt - assert db_service.db_file.visible_health_status == FileSystemItemHealthStatus.GOOD # not scanned yet + assert db_service.db_file.visible_health_status == FileSystemItemHealthStatus.NONE # not scanned yet db_service.db_file.scan() # scan the db file @@ -232,7 +232,7 @@ def test_database_service_fix(uc2_network): assert db_service.health_state_actual == SoftwareHealthState.FIXING # apply timestep until the fix is applied - for i in range(db_service.fixing_duration + 1): + for i in range(db_service.config.fixing_duration + 1): uc2_network.apply_timestep(i) assert db_service.db_file.health_status == FileSystemItemHealthStatus.GOOD @@ -266,7 +266,7 @@ def test_database_cannot_be_queried_while_fixing(uc2_network): assert db_connection.query(sql="SELECT") is False # apply timestep until the fix is applied - for i in range(db_service.fixing_duration + 1): + for i in range(db_service.config.fixing_duration + 1): uc2_network.apply_timestep(i) assert db_service.health_state_actual == SoftwareHealthState.GOOD @@ -308,7 +308,7 @@ def test_database_can_create_connection_while_fixing(uc2_network): assert new_db_connection.query(sql="SELECT") is False # still should fail to query because FIXING # apply timestep until the fix is applied - for i in range(db_service.fixing_duration + 1): + for i in range(db_service.config.fixing_duration + 1): uc2_network.apply_timestep(i) assert db_service.health_state_actual == SoftwareHealthState.GOOD diff --git a/tests/integration_tests/system/test_service_listening_on_ports.py b/tests/integration_tests/system/test_service_listening_on_ports.py index 2d3679ed..84413ac9 100644 --- a/tests/integration_tests/system/test_service_listening_on_ports.py +++ b/tests/integration_tests/system/test_service_listening_on_ports.py @@ -14,7 +14,14 @@ from primaite.utils.validation.port import PORT_LOOKUP from tests import TEST_ASSETS_ROOT -class _DatabaseListener(Service): +class _DatabaseListener(Service, identifier="_DatabaseListener"): + class ConfigSchema(Service.ConfigSchema): + """ConfigSchema for _DatabaseListener.""" + + type: str = "_DatabaseListener" + listen_on_ports: Set[int] = {PORT_LOOKUP["POSTGRES_SERVER"]} + + config: "_DatabaseListener.ConfigSchema" = Field(default_factory=lambda: _DatabaseListener.ConfigSchema()) name: str = "DatabaseListener" protocol: str = PROTOCOL_LOOKUP["TCP"] port: int = PORT_LOOKUP["NONE"] diff --git a/tests/integration_tests/system/test_web_client_server.py b/tests/integration_tests/system/test_web_client_server.py index c1028e8e..8aea34c1 100644 --- a/tests/integration_tests/system/test_web_client_server.py +++ b/tests/integration_tests/system/test_web_client_server.py @@ -51,7 +51,7 @@ def test_web_page_get_users_page_request_with_domain_name(web_client_and_web_ser web_browser_app, computer, web_server_service, server = web_client_and_web_server web_server_ip = server.network_interfaces.get(next(iter(server.network_interfaces))).ip_address - web_browser_app.target_url = f"http://arcd.com/" + web_browser_app.config.target_url = f"http://arcd.com/" assert web_browser_app.operating_state == ApplicationOperatingState.RUNNING assert web_browser_app.get_webpage() is True @@ -66,7 +66,7 @@ def test_web_page_get_users_page_request_with_ip_address(web_client_and_web_serv web_browser_app, computer, web_server_service, server = web_client_and_web_server web_server_ip = server.network_interfaces.get(next(iter(server.network_interfaces))).ip_address - web_browser_app.target_url = f"http://{web_server_ip}/" + web_browser_app.config.target_url = f"http://{web_server_ip}/" assert web_browser_app.operating_state == ApplicationOperatingState.RUNNING assert web_browser_app.get_webpage() is True @@ -81,7 +81,7 @@ def test_web_page_request_from_shut_down_server(web_client_and_web_server): web_browser_app, computer, web_server_service, server = web_client_and_web_server web_server_ip = server.network_interfaces.get(next(iter(server.network_interfaces))).ip_address - web_browser_app.target_url = f"http://arcd.com/" + web_browser_app.config.target_url = f"http://arcd.com/" assert web_browser_app.operating_state == ApplicationOperatingState.RUNNING assert web_browser_app.get_webpage() is True @@ -108,7 +108,7 @@ def test_web_page_request_from_closed_web_browser(web_client_and_web_server): web_browser_app, computer, web_server_service, server = web_client_and_web_server assert web_browser_app.operating_state == ApplicationOperatingState.RUNNING - web_browser_app.target_url = f"http://arcd.com/" + web_browser_app.config.target_url = f"http://arcd.com/" assert web_browser_app.get_webpage() is True # latest response should have status code 200 diff --git a/tests/integration_tests/system/test_web_client_server_and_database.py b/tests/integration_tests/system/test_web_client_server_and_database.py index 8fb6dc18..41f1a837 100644 --- a/tests/integration_tests/system/test_web_client_server_and_database.py +++ b/tests/integration_tests/system/test_web_client_server_and_database.py @@ -74,7 +74,7 @@ def web_client_web_server_database(example_network) -> Tuple[Network, Computer, # Install Web Browser on computer computer.software_manager.install(WebBrowser) web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser") - web_browser.target_url = "http://arcd.com/users/" + web_browser.config.target_url = "http://arcd.com/users/" web_browser.run() # Install DNS Client service on computer @@ -131,7 +131,7 @@ def test_database_fix_disrupts_web_client(uc2_network): assert web_browser.get_webpage() is False - for i in range(database_service.fixing_duration + 1): + for i in range(database_service.config.fixing_duration + 1): uc2_network.apply_timestep(i) assert database_service.health_state_actual == SoftwareHealthState.GOOD diff --git a/tests/unit_tests/_primaite/_game/_agent/test_actions.py b/tests/unit_tests/_primaite/_game/_agent/test_actions.py index 63ea1b07..5750befd 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_actions.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_actions.py @@ -2,10 +2,11 @@ from unittest.mock import Mock import pytest +from pydantic import ValidationError -from primaite.game.agent.actions import ( - ActionManager, - DoNothingAction, +from primaite.game.agent.actions import ActionManager +from primaite.game.agent.actions.manager import DoNothingAction +from primaite.game.agent.actions.service import ( NodeServiceDisableAction, NodeServiceEnableAction, NodeServicePauseAction, @@ -18,13 +19,8 @@ from primaite.game.agent.actions import ( def test_do_nothing_action_form_request(): - """Test that the DoNothingAction can form a request and that it is correct.""" - manager = Mock() - - action = DoNothingAction(manager=manager) - - request = action.form_request() - + """Test that the do_nothingAction can form a request and that it is correct.""" + request = DoNothingAction.form_request(DoNothingAction.ConfigSchema()) assert request == ["do_nothing"] @@ -42,7 +38,7 @@ def test_do_nothing_action_form_request(): ], ) # flake8: noqa @pytest.mark.parametrize( - "node_name, service_name, expect_to_do_nothing", + "node_name, service_name, expect_failure", [ ("pc_1", "chrome", False), (None, "chrome", True), @@ -50,42 +46,15 @@ def test_do_nothing_action_form_request(): (None, None, True), ], ) # flake8: noqa -def test_service_action_form_request(node_name, service_name, expect_to_do_nothing, action_class, action_verb): +def test_service_action_form_request(node_name, service_name, expect_failure, action_class, action_verb): """Test that the ServiceScanAction can form a request and that it is correct.""" - manager: ActionManager = Mock() - manager.get_node_name_by_idx.return_value = node_name - manager.get_service_name_by_idx.return_value = service_name - - action = action_class(manager=manager, num_nodes=1, num_services=1) - - request = action.form_request(node_id=0, service_id=0) - - if expect_to_do_nothing: - assert request == ["do_nothing"] + if expect_failure: + with pytest.raises(ValidationError): + request = action_class.form_request( + config=action_class.ConfigSchema(node_name=node_name, service_name=service_name) + ) else: + request = action_class.form_request( + config=action_class.ConfigSchema(node_name=node_name, service_name=service_name) + ) assert request == ["network", "node", node_name, "service", service_name, action_verb] - - -@pytest.mark.parametrize( - "node_name, service_name, expect_to_do_nothing", - [ - ("pc_1", "chrome", False), - (None, "chrome", True), - ("pc_1", None, True), - (None, None, True), - ], -) # flake8: noqa -def test_service_scan_form_request(node_name, service_name, expect_to_do_nothing): - """Test that the ServiceScanAction can form a request and that it is correct.""" - manager: ActionManager = Mock() - manager.get_node_name_by_idx.return_value = node_name - manager.get_service_name_by_idx.return_value = service_name - - action = NodeServiceScanAction(manager=manager, num_nodes=1, num_services=1) - - request = action.form_request(node_id=0, service_id=0) - - if expect_to_do_nothing: - assert request == ["do_nothing"] - else: - assert request == ["network", "node", node_name, "service", service_name, "scan"] diff --git a/tests/unit_tests/_primaite/_game/_agent/test_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_agent.py new file mode 100644 index 00000000..7956a44f --- /dev/null +++ b/tests/unit_tests/_primaite/_game/_agent/test_agent.py @@ -0,0 +1,52 @@ +from primaite.game.agent.observations.file_system_observations import FileObservation +from primaite.game.agent.observations.observation_manager import NullObservation +from primaite.game.agent.scripted_agents.random_agent import RandomAgent + + +def test_creating_empty_agent(): + agent = RandomAgent() + assert len(agent.action_manager.action_map) == 0 + assert isinstance(agent.observation_manager.obs, NullObservation) + assert len(agent.reward_function.reward_components) == 0 + + +def test_creating_agent_from_dict(): + action_config = { + "action_map": { + 0: {"action": "do_nothing", "options": {}}, + 1: { + "action": "node_application_execute", + "options": {"node_name": "client", "application_name": "database"}, + }, + } + } + observation_config = { + "type": "FILE", + "options": { + "file_name": "dog.pdf", + "include_num_access": False, + "file_system_requires_scan": False, + }, + } + reward_config = { + "reward_components": [ + { + "type": "DATABASE_FILE_INTEGRITY", + "weight": 0.3, + "options": {"node_hostname": "server", "folder_name": "database", "file_name": "database.db"}, + } + ] + } + agent = RandomAgent( + config={ + "ref": "random_agent", + "team": "BLUE", + "action_space": action_config, + "observation_space": observation_config, + "reward_function": reward_config, + } + ) + + assert len(agent.action_manager.action_map) == 2 + assert isinstance(agent.observation_manager.obs, FileObservation) + assert len(agent.reward_function.reward_components) == 1 diff --git a/tests/unit_tests/_primaite/_game/_agent/test_observations.py b/tests/unit_tests/_primaite/_game/_agent/test_observations.py index bb3ad33c..5156a29f 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_observations.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_observations.py @@ -69,8 +69,8 @@ class TestFileSystemRequiresScan: wildcard_list: - 0.0.0.1 port_list: - - 80 - - 5432 + - HTTP + - POSTGRES_SERVER protocol_list: - ICMP - TCP @@ -98,7 +98,7 @@ class TestFileSystemRequiresScan: """ cfg = yaml.safe_load(obs_cfg_yaml) - manager = ObservationManager.from_config(cfg) + manager = ObservationManager(config=cfg) hosts: List[HostObservation] = manager.obs.components["NODES"].hosts for i, host in enumerate(hosts): @@ -119,14 +119,20 @@ class TestFileSystemRequiresScan: assert obs_not_requiring_scan.observe(file_state)["health_status"] == 3 def test_folder_require_scan(self): - folder_state = {"health_status": 3, "visible_status": 1} + folder_state = {"health_status": 3, "visible_status": 1, "scanned_this_step": False} obs_requiring_scan = FolderObservation( [], files=[], num_files=0, include_num_access=False, file_system_requires_scan=True ) - assert obs_requiring_scan.observe(folder_state)["health_status"] == 1 + assert obs_requiring_scan.observe(folder_state)["health_status"] == 0 obs_not_requiring_scan = FolderObservation( [], files=[], num_files=0, include_num_access=False, file_system_requires_scan=False ) assert obs_not_requiring_scan.observe(folder_state)["health_status"] == 3 + + folder_state = {"health_status": 3, "visible_status": 1, "scanned_this_step": True} + obs_requiring_scan = FolderObservation( + [], files=[], num_files=0, include_num_access=False, file_system_requires_scan=True + ) + assert obs_requiring_scan.observe(folder_state)["health_status"] == 1 diff --git a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py index d8f800e4..f55033fd 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py @@ -3,6 +3,7 @@ from primaite.game.agent.actions import ActionManager from primaite.game.agent.observations.observation_manager import NestedObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent +from primaite.game.game import PrimaiteGame, PrimaiteGameOptions def test_probabilistic_agent(): @@ -16,69 +17,58 @@ def test_probabilistic_agent(): """ N_TRIALS = 10_000 P_DO_NOTHING = 0.1 - P_NODE_APPLICATION_EXECUTE = 0.3 - P_NODE_FILE_DELETE = 0.6 + P_node_application_execute = 0.3 + P_node_file_delete = 0.6 MIN_DO_NOTHING = 850 MAX_DO_NOTHING = 1150 - MIN_NODE_APPLICATION_EXECUTE = 2800 - MAX_NODE_APPLICATION_EXECUTE = 3200 - MIN_NODE_FILE_DELETE = 5750 - MAX_NODE_FILE_DELETE = 6250 + MIN_node_application_execute = 2800 + MAX_node_application_execute = 3200 + MIN_node_file_delete = 5750 + MAX_node_file_delete = 6250 - action_space = ActionManager( - actions=[ - {"type": "DONOTHING"}, - {"type": "NODE_APPLICATION_EXECUTE"}, - {"type": "NODE_FILE_DELETE"}, - ], - nodes=[ - { - "node_name": "client_1", - "applications": [{"application_name": "WebBrowser"}], - "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}], + action_space_cfg = { + "action_map": { + 0: {"action": "do_nothing", "options": {}}, + 1: { + "action": "node_application_execute", + "options": {"node_name": "client_1", "application_name": "WebBrowser"}, + }, + 2: { + "action": "node_file_delete", + "options": {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, }, - ], - max_folders_per_node=2, - max_files_per_folder=2, - max_services_per_node=2, - max_applications_per_node=2, - max_nics_per_node=2, - max_acl_rules=10, - protocols=["TCP", "UDP", "ICMP"], - ports=["HTTP", "DNS", "ARP"], - act_map={ - 0: {"action": "DONOTHING", "options": {}}, - 1: {"action": "NODE_APPLICATION_EXECUTE", "options": {"node_id": 0, "application_id": 0}}, - 2: {"action": "NODE_FILE_DELETE", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}}, }, - ) - observation_space = ObservationManager(NestedObservation(components={})) - reward_function = RewardFunction() + } - pa = ProbabilisticAgent( - agent_name="test_agent", - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - settings={ - "action_probabilities": {0: P_DO_NOTHING, 1: P_NODE_APPLICATION_EXECUTE, 2: P_NODE_FILE_DELETE}, + game = PrimaiteGame() + game.options = PrimaiteGameOptions(ports=[], protocols=[]) + + pa_config = { + "type": "ProbabilisticAgent", + "ref": "ProbabilisticAgent", + "team": "BLUE", + "action_space": action_space_cfg, + "agent_settings": { + "action_probabilities": {0: P_DO_NOTHING, 1: P_node_application_execute, 2: P_node_file_delete}, }, - ) + } + + pa = ProbabilisticAgent.from_config(config=pa_config) do_nothing_count = 0 node_application_execute_count = 0 node_file_delete_count = 0 for _ in range(N_TRIALS): a = pa.get_action(0) - if a == ("DONOTHING", {}): + if a == ("do_nothing", {}): do_nothing_count += 1 - elif a == ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}): + elif a == ("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"}): node_application_execute_count += 1 - elif a == ("NODE_FILE_DELETE", {"node_id": 0, "folder_id": 0, "file_id": 0}): + elif a == ("node_file_delete", {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}): node_file_delete_count += 1 else: raise AssertionError("Probabilistic agent produced an unexpected action.") assert MIN_DO_NOTHING < do_nothing_count < MAX_DO_NOTHING - assert MIN_NODE_APPLICATION_EXECUTE < node_application_execute_count < MAX_NODE_APPLICATION_EXECUTE - assert MIN_NODE_FILE_DELETE < node_file_delete_count < MAX_NODE_FILE_DELETE + assert MIN_node_application_execute < node_application_execute_count < MAX_node_application_execute + assert MIN_node_file_delete < node_file_delete_count < MAX_node_file_delete diff --git a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py index 346eb512..5af71319 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py @@ -81,7 +81,7 @@ class TestWebpageUnavailabilitySticky: reward = WebpageUnavailablePenalty(config=schema) # no response codes yet, reward is 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) browser_history = [] state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} @@ -91,8 +91,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == 0 # agent did a successful fetch - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="success", data={}) browser_history.append({"outcome": 200}) @@ -104,7 +104,7 @@ class TestWebpageUnavailabilitySticky: # THE IMPORTANT BIT # agent did nothing, because reward is not sticky, it goes back to 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) browser_history = [] state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} @@ -114,8 +114,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == 0.0 # agent fails to fetch, get a -1.0 reward - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -126,8 +126,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == -1.0 # agent fails again to fetch, get a -1.0 reward again - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -142,7 +142,7 @@ class TestWebpageUnavailabilitySticky: reward = WebpageUnavailablePenalty(config=schema) # no response codes yet, reward is 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) browser_history = [] state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} @@ -152,8 +152,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == 0 # agent did a successful fetch - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="success", data={}) browser_history.append({"outcome": 200}) @@ -165,7 +165,7 @@ class TestWebpageUnavailabilitySticky: # THE IMPORTANT BIT # agent did nothing, because reward is sticky, it stays at 1.0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} last_action_response = AgentHistoryItem( @@ -174,8 +174,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == 1.0 # agent fails to fetch, get a -1.0 reward - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -186,8 +186,8 @@ class TestWebpageUnavailabilitySticky: assert reward.calculate(state, last_action_response) == -1.0 # agent fails again to fetch, get a -1.0 reward again - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -207,7 +207,7 @@ class TestGreenAdminDatabaseUnreachableSticky: reward = GreenAdminDatabaseUnreachablePenalty(config=schema) # no response codes yet, reward is 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} last_action_response = AgentHistoryItem( @@ -216,8 +216,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == 0 # agent did a successful fetch - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -228,7 +228,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # THE IMPORTANT BIT # agent did nothing, because reward is not sticky, it goes back to 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} last_action_response = AgentHistoryItem( @@ -237,8 +237,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == 0.0 # agent fails to fetch, get a -1.0 reward - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -248,8 +248,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == -1.0 # agent fails again to fetch, get a -1.0 reward again - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -266,7 +266,7 @@ class TestGreenAdminDatabaseUnreachableSticky: reward = GreenAdminDatabaseUnreachablePenalty(config=schema) # no response codes yet, reward is 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} last_action_response = AgentHistoryItem( @@ -275,8 +275,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == 0 # agent did a successful fetch - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -287,7 +287,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # THE IMPORTANT BIT # agent did nothing, because reward is not sticky, it goes back to 0 - action, params, request = "DO_NOTHING", {}, ["DONOTHING"] + action, params, request = "do_nothing", {}, ["do_nothing"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} last_action_response = AgentHistoryItem( @@ -296,8 +296,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == 1.0 # agent fails to fetch, get a -1.0 reward - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -307,8 +307,8 @@ class TestGreenAdminDatabaseUnreachableSticky: assert reward.calculate(state, last_action_response) == -1.0 # agent fails again to fetch, get a -1.0 reward again - action = "NODE_APPLICATION_EXECUTE" - params = {"node_id": 0, "application_id": 0} + action = "node_application_execute" + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py index 6cbf93c8..9691080d 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py @@ -22,12 +22,12 @@ def test_file_scan(file_system): file: File = file_system.create_file(file_name="test_file.txt", folder_name="test_folder") assert file.health_status == FileSystemItemHealthStatus.GOOD - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE file.corrupt() assert file.health_status == FileSystemItemHealthStatus.CORRUPT - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE file.scan() @@ -46,7 +46,7 @@ def test_file_reveal_to_red_scan(file_system): assert file.revealed_to_red is True -@pytest.mark.skip(reason="NODE_FILE_CHECKHASH not implemented") +@pytest.mark.skip(reason="node_file_checkhash not implemented") def test_simulated_file_check_hash(file_system): file: File = file_system.create_file(file_name="test_file.txt", folder_name="test_folder") diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py index 4ec1ec57..59f3f000 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py @@ -24,7 +24,7 @@ def test_file_scan_request(populated_file_system): file.corrupt() assert file.health_status == FileSystemItemHealthStatus.CORRUPT - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE fs.apply_request(request=["folder", folder.name, "file", file.name, "scan"]) @@ -32,7 +32,7 @@ def test_file_scan_request(populated_file_system): assert file.visible_health_status == FileSystemItemHealthStatus.CORRUPT -@pytest.mark.skip(reason="NODE_FILE_CHECKHASH not implemented") +@pytest.mark.skip(reason="node_file_checkhash not implemented") def test_file_checkhash_request(populated_file_system): """Test that an agent can request a file hash check.""" fs, folder, file = populated_file_system @@ -94,7 +94,7 @@ def test_deleted_file_cannot_be_interacted_with(populated_file_system): assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert ( fs.get_file(folder_name=folder.name, file_name=file.name).visible_health_status - == FileSystemItemHealthStatus.GOOD + == FileSystemItemHealthStatus.NONE ) fs.apply_request(request=["delete", "file", folder.name, file.name]) diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py index 473e0db2..b5d9b269 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py @@ -44,25 +44,25 @@ def test_folder_scan(file_system): file2: File = folder.get_file_by_id(file_uuid=list(folder.files)[0]) assert folder.health_status == FileSystemItemHealthStatus.GOOD - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE + assert file1.visible_health_status == FileSystemItemHealthStatus.NONE + assert file2.visible_health_status == FileSystemItemHealthStatus.NONE folder.corrupt() assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE + assert file1.visible_health_status == FileSystemItemHealthStatus.NONE + assert file2.visible_health_status == FileSystemItemHealthStatus.NONE folder.scan() folder.apply_timestep(timestep=0) assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE + assert file1.visible_health_status == FileSystemItemHealthStatus.NONE + assert file2.visible_health_status == FileSystemItemHealthStatus.NONE folder.apply_timestep(timestep=1) folder.apply_timestep(timestep=2) @@ -120,7 +120,7 @@ def test_folder_corrupt_repair(file_system): assert file.health_status == FileSystemItemHealthStatus.GOOD -@pytest.mark.skip(reason="NODE_FILE_CHECKHASH not implemented") +@pytest.mark.skip(reason="node_file_checkhash not implemented") def test_simulated_folder_check_hash(file_system): folder: Folder = file_system.create_folder(folder_name="test_folder") file_system.create_file(file_name="test_file.txt", folder_name="test_folder") diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py index 609e29c4..72857638 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py @@ -29,18 +29,18 @@ def test_folder_scan_request(populated_file_system): folder.corrupt() assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE + assert file1.visible_health_status == FileSystemItemHealthStatus.NONE + assert file2.visible_health_status == FileSystemItemHealthStatus.NONE fs.apply_request(request=["folder", folder.name, "scan"]) folder.apply_timestep(timestep=0) assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD - assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE + assert file1.visible_health_status == FileSystemItemHealthStatus.NONE + assert file2.visible_health_status == FileSystemItemHealthStatus.NONE folder.apply_timestep(timestep=1) folder.apply_timestep(timestep=2) @@ -51,7 +51,7 @@ def test_folder_scan_request(populated_file_system): assert file2.visible_health_status == FileSystemItemHealthStatus.CORRUPT -@pytest.mark.skip(reason="NODE_FOLDER_CHECKHASH not implemented") +@pytest.mark.skip(reason="node_folder_checkhash not implemented") def test_folder_checkhash_request(populated_file_system): """Test that an agent can request a folder hash check.""" fs, folder, file = populated_file_system diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py index 605f8c3b..672a4b5f 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py @@ -70,13 +70,13 @@ def test_node_os_scan(node): # add folder and file to node folder: Folder = node.file_system.create_folder(folder_name="test_folder") folder.corrupt() - assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD + assert folder.visible_health_status == FileSystemItemHealthStatus.NONE file: File = node.file_system.create_file(folder_name="test_folder", file_name="file.txt") file2: File = node.file_system.create_file(folder_name="test_folder", file_name="file2.txt") file.corrupt() file2.corrupt() - assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.NONE # run os scan node.apply_request(["os", "scan"]) diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py index 4ff387ce..17f8445a 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py @@ -128,13 +128,13 @@ def test_c2_handle_switching_port(basic_c2_network): assert c2_server.c2_connection_active is True # Assert to confirm that both the C2 server and the C2 beacon are configured correctly. - assert c2_beacon.c2_config.keep_alive_frequency is 2 - assert c2_beacon.c2_config.masquerade_port is PORT_LOOKUP["HTTP"] - assert c2_beacon.c2_config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] + assert c2_beacon.config.keep_alive_frequency is 2 + assert c2_beacon.config.masquerade_port is PORT_LOOKUP["HTTP"] + assert c2_beacon.config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] - assert c2_server.c2_config.keep_alive_frequency is 2 - assert c2_server.c2_config.masquerade_port is PORT_LOOKUP["HTTP"] - assert c2_server.c2_config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] + assert c2_server.config.keep_alive_frequency is 2 + assert c2_server.config.masquerade_port is PORT_LOOKUP["HTTP"] + assert c2_server.config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] # Configuring the C2 Beacon. c2_beacon.configure( @@ -150,11 +150,11 @@ def test_c2_handle_switching_port(basic_c2_network): # Assert to confirm that both the C2 server and the C2 beacon # Have reconfigured their C2 settings. - assert c2_beacon.c2_config.masquerade_port is PORT_LOOKUP["FTP"] - assert c2_beacon.c2_config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] + assert c2_beacon.config.masquerade_port is PORT_LOOKUP["FTP"] + assert c2_beacon.config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] - assert c2_server.c2_config.masquerade_port is PORT_LOOKUP["FTP"] - assert c2_server.c2_config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] + assert c2_server.config.masquerade_port is PORT_LOOKUP["FTP"] + assert c2_server.config.masquerade_protocol is PROTOCOL_LOOKUP["TCP"] def test_c2_handle_switching_frequency(basic_c2_network): @@ -174,8 +174,8 @@ def test_c2_handle_switching_frequency(basic_c2_network): assert c2_server.c2_connection_active is True # Assert to confirm that both the C2 server and the C2 beacon are configured correctly. - assert c2_beacon.c2_config.keep_alive_frequency is 2 - assert c2_server.c2_config.keep_alive_frequency is 2 + assert c2_beacon.config.keep_alive_frequency is 2 + assert c2_server.config.keep_alive_frequency is 2 # Configuring the C2 Beacon. c2_beacon.configure(c2_server_ip_address="192.168.0.1", keep_alive_frequency=10) @@ -186,8 +186,8 @@ def test_c2_handle_switching_frequency(basic_c2_network): # Assert to confirm that both the C2 server and the C2 beacon # Have reconfigured their C2 settings. - assert c2_beacon.c2_config.keep_alive_frequency is 10 - assert c2_server.c2_config.keep_alive_frequency is 10 + assert c2_beacon.config.keep_alive_frequency is 10 + assert c2_server.config.keep_alive_frequency is 10 # Now skipping 9 time steps to confirm keep alive inactivity for i in range(9): diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py index ad6fe135..5598e1a7 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py @@ -148,7 +148,7 @@ def test_service_fixing(service): service.fix() assert service.health_state_actual == SoftwareHealthState.FIXING - for i in range(service.fixing_duration + 1): + for i in range(service.config.fixing_duration + 1): service.apply_timestep(i) assert service.health_state_actual == SoftwareHealthState.GOOD diff --git a/tests/unit_tests/_primaite/_simulator/_system/test_software.py b/tests/unit_tests/_primaite/_simulator/_system/test_software.py index a203a636..bdf9cfee 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/test_software.py +++ b/tests/unit_tests/_primaite/_simulator/_system/test_software.py @@ -2,6 +2,7 @@ from typing import Dict import pytest +from pydantic import Field from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.services.service import Service @@ -10,7 +11,14 @@ from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP -class TestSoftware(Service): +class TestSoftware(Service, identifier="TestSoftware"): + class ConfigSchema(Service.ConfigSchema): + """ConfigSChema for TestSoftware.""" + + type: str = "TestSoftware" + + config: "TestSoftware.ConfigSchema" = Field(default_factory=lambda: TestSoftware.ConfigSchema()) + def describe_state(self) -> Dict: pass