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

Merged PR 597: Standardise discriminator (f.k.a. identifiers) naming convention.

Browse Source 
## Summary
I mean I don't expect reviewers to go through every single change. It was painful enough to make all the changes myself, I had to write a script and regex the $*!+ out of the codebase to get this change through

## Test process
Tests pass. Searching through the codebase with regex queries to check for compliant strings.

## Checklist
- [ ] PR is linked to a **work item**
- [ ] **acceptance criteria** of linked ticket are met
- [ ] performed **self-review** of the code
- [ ] written **tests** for any new functionality added with this PR
- [ ] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [ ] updated the **change log**
- [ ] ran **pre-commit** checks for code style
- [ ] attended to any **TO-DOs** left in the code

Related work items: #3062
This commit is contained in:
Marek Wolan
2025-02-05 10:49:40 +00:00
parent 0708c836a3 4a472c5c75
commit 01bf052cc4
208 changed files with 2952 additions and 2948 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
docs/source/action_masking.rst
View File

@@ -23,117 +23,117 @@ The following logic is applied:
+------------------------------------------+---------------------------------------------------------------------+
| Action | Action Mask Logic |
+==========================================+=====================================================================+
| **do_nothing** | 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_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-add-rule** | Router is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **router_acl_removerule** | Router is on. |
| **router-acl-remove-rule** | Router is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **firewall_acl_addrule** | Firewall is on. |
| **firewall-acl-add-rule** | Firewall is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **firewall_acl_removerule** | Firewall is on. |
| **firewall-acl-remove-rule** | Firewall is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **configure_database_client** | Node is on. |
| **configure-database-client** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **configure_ransomware_script** | Node is on. |
| **configure-ransomware-script** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **c2_server_ransomware_configure** | Node is on. |
| **c2-server-ransomware-configure** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **configure_dos_bot** | Node is on. |
| **configure-dos-bot** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **configure_c2_beacon** | Node is on. |
| **configure-c2-beacon** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **c2_server_ransomware_launch** | Node is on. |
| **c2-server-ransomware-launch** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **c2_server_terminal_command** | Node is on. |
| **c2-server-terminal-command** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **c2_server_data_exfiltrate** | Node is on. |
| **c2-server-data-exfiltrate** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **node_account_change_password** | Node is on. |
| **node-account-change-password** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **node_session_remote_login** | Node is on. |
| **node-session-remote-login** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **node_session_remote_logoff** | Node is on. |
| **node-session-remote-logoff** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+
| **node_send_remote_command** | Node is on. |
| **node-send-remote-command** | Node is on. |
+------------------------------------------+---------------------------------------------------------------------+

22
docs/source/configuration/agents.rst
View File

@@ -19,13 +19,13 @@ Agents can be scripted (deterministic and stochastic), or controlled by a reinfo
...
- ref: green_agent_example
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
observation_space:
type: UC2GreenObservation
action_space:
reward_function:
reward_components:
- type: DUMMY
- type: dummy
agent_settings:
start_settings:
@@ -44,13 +44,13 @@ Specifies if the agent is malicious (``RED``), benign (``GREEN``), or defensive
``type``
--------
Specifies which class should be used for the agent. ``ProxyAgent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``RedDatabaseCorruptingAgent`` and ``ProbabilisticAgent`` generate their own behaviour.
Specifies which class should be used for the agent. ``proxy-agent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``red-database-corrupting-agent`` and ``probabilistic-agent`` generate their own behaviour.
Available agent types:
- ``ProbabilisticAgent``
- ``ProxyAgent``
- ``RedDatabaseCorruptingAgent``
- ``probabilistic-agent``
- ``proxy-agent``
- ``red-database-corrupting-agent``
``observation_space``
---------------------
@@ -66,10 +66,10 @@ selects which python class from the :py:mod:`primaite.game.agent.observation` mo
Allows configuration of the chosen observation type. These are optional.
* ``num_services_per_node``, ``num_folders_per_node``, ``num_files_per_folder``, ``num_nics_per_node`` all define the shape of the observation space. The size and shape of the obs space must remain constant, but the number of files, folders, ACL rules, and other components can change within an episode. Therefore padding is performed and these options set the size of the obs space.
* ``num_services_per_node``, ``num_folders_per_node``, ``num_files_per_folder``, ``num_nics_per_node`` all define the shape of the observation space. The size and shape of the obs space must remain constant, but the number of files, folders, acl rules, and other components can change within an episode. Therefore padding is performed and these options set the size of the obs space.
* ``nodes``: list of nodes that will be present in this agent's observation space. The ``node_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config. Each node can also be configured with services, and files that should be monitored.
* ``links``: list of links that will be present in this agent's observation space. The ``link_ref`` relates to the human-readable unique reference defined later in the ``simulation`` part of the config.
* ``acl``: configure how the agent reads the access control list on the router in the simulation. ``router_node_ref`` is for selecting which router's ACL table should be used. ``ip_list`` sets the encoding of ip addresses as integers within the observation space.
* ``acl``: configure how the agent reads the access control list on the router in the simulation. ``router_node_ref`` is for selecting which router's acl table should be used. ``ip_list`` sets the encoding of ip addresses as integers within the observation space.
For more information see :py:mod:`primaite.game.agent.observations`
@@ -103,7 +103,7 @@ Similar to action space, this is defined as a list of components from the :py:mo
``reward_components``
^^^^^^^^^^^^^^^^^^^^^
TODO: update description
A list of reward types from :py:mod:`primaite.game.agent.rewards.RewardFunction.rew_class_identifiers`
e.g.
@@ -111,8 +111,8 @@ e.g.
.. code-block:: yaml
reward_components:
- type: DUMMY
- type: DATABASE_FILE_INTEGRITY
- type: dummy
- type: database-file-integrity
``agent_settings``

2
docs/source/configuration/simulation.rst
View File

@@ -6,7 +6,7 @@
``simulation``
==============
In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents.
# TODO: ref field is no longer real
At the top level of the network are ``nodes``, ``links`` and ``airspace``.
e.g.

32
docs/source/configuration/simulation/nodes/network_examples.rst
View File

@@ -617,10 +617,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai
@@ -631,10 +631,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai
@@ -700,7 +700,7 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 8.8.8.1
services:
- ref: dns_server
type: DNSServer
type: dns-server
options:
domain_mapping:
sometech.ai: 94.10.180.6
@@ -794,9 +794,9 @@ Each node is configured to ensure it meets the specific security and operational
dns_server: 8.8.8.2
services:
- ref: web_server
type: WebServer
type: web-server
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
@@ -903,10 +903,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: DatabaseService
- type: database-service
options:
backup_server_ip: 10.10.1.12 # The some_tech_storage_srv server
- type: FTPClient
- type: ftp-client
- hostname: some_tech_storage_srv
type: server
@@ -915,7 +915,7 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: FTPServer
- type: ftp-server
- hostname: some_tech_hr_1
type: computer
@@ -924,10 +924,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 10.10.3.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai
@@ -938,10 +938,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai
@@ -952,10 +952,10 @@ Each node is configured to ensure it meets the specific security and operational
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai

8
docs/source/game_layer.rst
View File

@@ -57,13 +57,13 @@ An agent's reward can be based on rewards of other agents. This is particularly
reward_components:
# When the webpage loads, the reward goes up by 0.25 when it fails to load, it goes down to -0.25
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
weight: 0.25
options:
node_hostname: client_2
# When the database is reachable, the reward goes up by 0.05, when it is unreachable it goes down to -0.05
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 0.05
options:
node_hostname: client_2
@@ -74,7 +74,7 @@ An agent's reward can be based on rewards of other agents. This is particularly
reward_components:
# When the database file is in a good state, blue's reward is 0.4, when it's in a corrupted state the reward is -0.4
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.40
options:
node_hostname: database_server
@@ -82,7 +82,7 @@ An agent's reward can be based on rewards of other agents. This is particularly
file_name: database.db
# The green's reward is added onto the blue's reward.
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: client_2_green_user

10
docs/source/how_to_guides/extensible_actions.rst
View File

@@ -20,7 +20,7 @@ Custom actions within PrimAITE must be a sub-class of `AbstractAction`, and cont
#. ConfigSchema class
#. Unique Identifier
#. Unique discriminator
#. `form_request` method.
@@ -31,14 +31,14 @@ 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
Unique discriminator
#################
When declaring a custom class, it must have a unique identifier string, that allows PrimAITE to generate the correct action when needed.
When declaring a custom class, it must have a unique discriminator string, that allows PrimAITE to generate the correct action when needed.
.. code:: Python
class CreateDirectoryAction(AbstractAction, identifier="node_folder_create")
class CreateDirectoryAction(AbstractAction, discriminator="node-folder-create")
config: CreateDirectoryAction.ConfigSchema
@@ -58,7 +58,7 @@ When declaring a custom class, it must have a unique identifier string, that all
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`.
The above action would fail pydantic validation as the discriminator "node-folder-create" is already used by the `NodeFolderCreateAction`, and would create a duplicate listing within `AbstractAction._registry`.
form_request method

10
docs/source/how_to_guides/extensible_agents.rst
View File

@@ -25,7 +25,7 @@ The core features that should be implemented in any new agent are detailed below
.. code-block:: python
class ExampleAgent(AbstractAgent, identifier = "ExampleAgent"):
class ExampleAgent(AbstractAgent, discriminator = "ExampleAgent"):
"""An example agent for demonstration purposes."""
config: "ExampleAgent.ConfigSchema" = Field(default_factory= lambda: ExampleAgent.ConfigSchema())
@@ -51,11 +51,11 @@ The core features that should be implemented in any new agent are detailed below
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
reward_function:
reward_components:
- type: DUMMY
- type: dummy
agent_settings:
start_step: 25
@@ -64,9 +64,9 @@ The core features that should be implemented in any new agent are detailed below
starting_host: "Server_1"
#. **Identifiers**:
#. **discriminators**:
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.
All agent classes should have an ``discriminator`` 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
====================

6
docs/source/how_to_guides/extensible_rewards.rst
View File

@@ -17,7 +17,7 @@ Reward classes are inherited from AbstractReward (a sub-class of Pydantic's Base
Within the reward class there is a ConfigSchema class responsible for ensuring the config file data
is in the correct format. This also means there is little (if no) requirement for and `__init__`
method. The `.from_config` method is no longer required as it's inherited from `AbstractReward`.
Each class requires an identifier string which is used by the ConfigSchema class to verify that it
Each class requires an discriminator string which is used by the ConfigSchema class to verify that it
hasn't previously been added to the registry.
Inheriting from `BaseModel` removes the need for an `__init__` method but means that object
@@ -28,7 +28,7 @@ To add a new reward class follow the example below. Note that the type attribute
.. code-block:: Python
class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"):
class DatabaseFileIntegrity(AbstractReward, discriminator="database-file-integrity"):
"""Reward function component which rewards the agent for maintaining the integrity of a database file."""
config: "DatabaseFileIntegrity.ConfigSchema"
@@ -38,7 +38,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for DatabaseFileIntegrity."""
type: str = "DATABASE_FILE_INTEGRITY"
type: str = "database-file-integrity"
node_hostname: str
folder_name: str
file_name: str

8
docs/source/node_sets.rst
View File

@@ -55,7 +55,7 @@ Via YAML Config
nodes:
# ... nodes go here
node_sets:
- type: office_lan
- type: office-lan
lan_name: CORP_LAN
subnet_base: 2
pcs_ip_block_start: 10
@@ -82,9 +82,9 @@ Here is an example of creating a custom node adder, DataCenterAdder:
.. code-block:: python
class DataCenterAdder(NetworkNodeAdder, identifier="data_center"):
class DataCenterAdder(NetworkNodeAdder, discriminator="data-center"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
type: Literal["data_center"] = "data_center"
type: Literal["data-center"] = "data-center"
num_servers: int
data_center_name: str
@@ -106,7 +106,7 @@ Here is an example of creating a custom node adder, DataCenterAdder:
.. code-block:: python
config = {
"type": "data_center",
"type": "data-center",
"num_servers": 5,
"data_center_name": "dc1"
}

4
docs/source/request_system.rst
View File

@@ -51,10 +51,10 @@ Request responses
When the simulator receives a request, it returns a response with a success status. The possible statuses are:
* **success**: The request was received and successfully executed.
* For example, the agent tries to add an ACL rule and specifies correct parameters, and the ACL rule is added successfully.
* For example, the agent tries to add an acl rule and specifies correct parameters, and the acl rule is added successfully.
* **failure**: The request was received, but it could not be executed, or it failed while executing.
* For example, the agent tries to execute the ``WebBrowser`` application, but the webpage wasn't retrieved because the DNS server is not setup on the node.
* For example, the agent tries to execute the ``web-browser`` application, but the webpage wasn't retrieved because the DNS server is not setup on the node.
* **unreachable**: The request was sent to a simulation component that does not exist.
* For example, the agent tries to scan a file that has not been created yet.

14
docs/source/rewards.rst
View File

@@ -23,7 +23,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: DUMMY
- type: dummy
weight: 1.0
@@ -36,7 +36,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 1.0
options:
node_hostname: server_1
@@ -53,7 +53,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: WEB_SERVER_404_PENALTY
- type: web-server-404-penalty
node_hostname: web_server
weight: 1.0
options:
@@ -70,7 +70,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
node_hostname: computer_1
weight: 1.0
options:
@@ -86,7 +86,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 1.0
options:
node_hostname: admin_pc_1
@@ -104,7 +104,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: scripted_agent
@@ -119,7 +119,7 @@ The following API pages describe the use of each reward component and the possib
# ...
reward_function:
reward_components:
- type: ACTION_PENALTY
- type: action-penalty
weight: 1.0
options:
action_penalty: -0.3

4
docs/source/simulation_components/system/applications/c2_suite.rst
View File

@@ -229,7 +229,7 @@ Via Configuration
type: computer
...
applications:
type: C2Server
type: c2-server
...
hostname: computer_b
type: computer
@@ -238,7 +238,7 @@ Via Configuration
# Either an agent must use application_execute.
# Or a if using the simulation layer - .establish().
applications:
type: C2Beacon
type: c2-beacon
options:
c2_server_ip_address: ...
keep_alive_frequency: 5

8
docs/source/simulation_components/system/applications/data_manipulation_bot.rst
View File

@@ -95,7 +95,7 @@ If not using the data manipulation bot manually, it needs to be used with a data
agents:
- ref: data_manipulation_red_bot
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
observation_space:
type: UC2RedObservation
@@ -115,7 +115,7 @@ If not using the data manipulation bot manually, it needs to be used with a data
action_space:
reward_function:
reward_components:
- type: DUMMY
- type: dummy
agent_settings:
start_settings:
@@ -132,14 +132,14 @@ If not using the data manipulation bot manually, it needs to be used with a data
# ... additional configuration here
applications:
- ref: data_manipulation_bot
type: DataManipulationBot
type: data-manipulation-bot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
- ref: web_server_database_client
type: DatabaseClient
type: database-client
options:
db_server_ip: 192.168.1.14

2
docs/source/simulation_components/system/applications/database_client.rst
View File

@@ -83,7 +83,7 @@ Via Configuration
...
applications:
- ref: database_client
type: DatabaseClient
type: database-client
options:
db_server_ip: 192.168.0.1

38
docs/source/simulation_components/system/applications/dos_bot.rst
View File

@@ -4,10 +4,10 @@
.. _DoSBot:
DoSBot
dos-bot
######
The ``DoSBot`` is an implementation of a Denial of Service attack within the PrimAITE simulation.
The ``dos-bot`` is an implementation of a Denial of Service attack within the PrimAITE simulation.
This specifically simulates a `Slow Loris attack`_.
.. _Slow Loris Attack: https://en.wikipedia.org/wiki/Slowloris_(computer_security)
@@ -15,20 +15,20 @@ This specifically simulates a `Slow Loris attack`_.
Key features
============
- Connects to the :ref:`DatabaseService` via the ``SoftwareManager``.
- Makes many connections to the :ref:`DatabaseService` which ends up using up the available connections.
- Connects to the :ref:`database-service` via the ``SoftwareManager``.
- Makes many connections to the :ref:`database-service` which ends up using up the available connections.
Usage
=====
- Configure with target IP address and optional password.
- use ``run`` to run the application_loop of DoSBot to begin attacks
- DoSBot runs through different actions at each timestep
- use ``run`` to run the application_loop of dos-bot to begin attacks
- dos-bot runs through different actions at each timestep
Implementation
==============
- Leverages :ref:`DatabaseClient` to create connections with :ref`DatabaseServer`.
- Leverages :ref:`database-client` to create connections with :ref`DatabaseServer`.
- Extends base Application class.
Examples
@@ -42,7 +42,7 @@ Python
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
from primaite.simulator.system.applications.red_applications.dos_bot import dos-bot
# Create Computer
computer = Computer(
@@ -54,11 +54,11 @@ Python
)
computer.power_on()
# Install DoSBot on computer
computer.software_manager.install(DoSBot)
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
# Install dos-bot on computer
computer.software_manager.install(dos-bot)
dos_bot: dos-bot = computer.software_manager.software.get("dos-bot")
# Configure the DoSBot
# Configure the dos-bot
dos_bot.configure(
target_ip_address=IPv4Address("192.168.0.10"),
payload="SPOOF DATA",
@@ -68,7 +68,7 @@ Python
max_sessions=1000
)
# run DoSBot
# run dos-bot
dos_bot.run()
@@ -86,7 +86,7 @@ Via Configuration
...
applications:
- ref: dos_bot
type: DoSBot
type: dos-bot
options:
target_ip_address: 192.168.0.10
payload: SPOOF DATA
@@ -101,7 +101,7 @@ Configuration
``target_ip_address``
"""""""""""""""""""""
IP address of the :ref:`DatabaseService` which the ``DataManipulationBot`` will try to attack.
IP address of the :ref:`database-service` which the ``data-manipulation-bot`` will try to attack.
This must be a valid octet i.e. in the range of ``0.0.0.0`` and ``255.255.255.255``.
@@ -119,7 +119,7 @@ See :ref:`List of IPProtocols <List of IPProtocols>` for a list of protocols.
Optional. Default value is ``None``.
The payload that the ``DoSBot`` sends as part of its attack.
The payload that the ``dos-bot`` sends as part of its attack.
.. include:: ../common/db_payload_list.rst
@@ -128,14 +128,14 @@ The payload that the ``DoSBot`` sends as part of its attack.
Optional. Default value is ``False``.
If ``True`` the ``DoSBot`` will maintain its attack.
If ``True`` the ``dos-bot`` will maintain its attack.
``port_scan_p_of_success``
""""""""""""""""""""""""""
Optional. Default value is ``0.1``.
The chance of the ``DoSBot`` to succeed with a port scan (and therefore continue the attack).
The chance of the ``dos-bot`` to succeed with a port scan (and therefore continue the attack).
This must be a float value between ``0`` and ``1``.
@@ -153,7 +153,7 @@ This must be a float value between ``0`` and ``1``.
Optional. Default value is ``1000``.
The maximum number of sessions the ``DoSBot`` is able to make.
The maximum number of sessions the ``dos-bot`` is able to make.
This must be an integer value equal to or greater than ``0``.

2
docs/source/simulation_components/system/applications/web_browser.rst
View File

@@ -85,7 +85,7 @@ Via Configuration
...
applications:
- ref: web_browser
type: WebBrowser
type: web-browser
options:
target_url: http://arcd.com/

2
docs/source/simulation_components/system/services/database_service.rst
View File

@@ -87,7 +87,7 @@ Via Configuration
...
services:
- ref: database_service
type: DatabaseService
type: database-service
options:
backup_server_ip: 192.168.0.10

2
docs/source/simulation_components/system/services/dns_client.rst
View File

@@ -77,7 +77,7 @@ Via Configuration
...
services:
- ref: dns_client
type: DNSClient
type: dns-client
options:
dns_server: 192.168.0.10

2
docs/source/simulation_components/system/services/dns_server.rst
View File

@@ -74,7 +74,7 @@ Via Configuration
...
services:
- ref: dns_server
type: DNSServer
type: dns-server
options:
domain_mapping:
arcd.com: 192.168.0.10

2
docs/source/simulation_components/system/services/ftp_client.rst
View File

@@ -78,7 +78,7 @@ Via Configuration
...
services:
- ref: ftp_client
type: FTPClient
type: ftp-client
Configuration
=============

2
docs/source/simulation_components/system/services/ftp_server.rst
View File

@@ -74,7 +74,7 @@ Via Configuration
...
services:
- ref: ftp_server
type: FTPServer
type: ftp-server
options:
server_password: test

2
docs/source/simulation_components/system/services/ntp_client.rst
View File

@@ -73,7 +73,7 @@ Via Configuration
...
services:
- ref: ntp_client
type: NTPClient
type: ntp-client
options:
ntp_server_ip: 192.168.0.10

2
docs/source/simulation_components/system/services/ntp_server.rst
View File

@@ -73,7 +73,7 @@ Via Configuration
...
services:
- ref: ntp_server
type: NTPServer
type: ntp-server
``Common Attributes``

2
docs/source/simulation_components/system/services/web_server.rst
View File

@@ -73,7 +73,7 @@ Via Configuration
...
services:
- ref: web_server
type: WebServer
type: web-server
``Common Attributes``

260
src/primaite/config/_package_data/data_manipulation.yaml
View File

@@ -24,7 +24,7 @@ game:
agents:
- ref: client_2_green_user
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.3
@@ -34,33 +34,33 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client_2
application_name: WebBrowser
application_name: web-browser
2:
action: node_application_execute
action: node-application-execute
options:
node_name: client_2
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
weight: 0.25
options:
node_hostname: client_2
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 0.05
options:
node_hostname: client_2
- ref: client_1_green_user
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.3
@@ -70,26 +70,26 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client_1
application_name: WebBrowser
application_name: web-browser
2:
action: node_application_execute
action: node-application-execute
options:
node_name: client_1
application_name: WebBrowser
application_name: web-browser
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
weight: 0.25
options:
node_hostname: client_1
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 0.05
options:
node_hostname: client_1
@@ -100,31 +100,31 @@ agents:
- ref: data_manipulation_attacker
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client_1, client_2]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 25
frequency: 20
variance: 5
- ref: defender
team: BLUE
type: ProxyAgent
type: proxy-agent
observation_space:
type: CUSTOM
type: custom
options:
components:
- type: NODES
- type: nodes
label: NODES
options:
hosts:
- hostname: domain_controller
- hostname: web_server
services:
- service_name: WebServer
- service_name: web-server
- hostname: database_server
folders:
- folder_name: database
@@ -169,7 +169,7 @@ agents:
- UDP
num_rules: 10
- type: LINKS
- type: links
label: LINKS
options:
link_references:
@@ -183,222 +183,222 @@ agents:
- switch_2:eth-1<->client_1:eth-1
- switch_2:eth-2<->client_2:eth-1
- switch_2:eth-7<->security_suite:eth-2
- type: "NONE"
- type: "none"
label: ICS
options: {}
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
# scan webapp service
1:
action: node_service_scan
action: node-service-scan
options:
node_name: web_server
service_name: WebServer
service_name: web-server
# stop webapp service
2:
action: node_service_stop
action: node-service-stop
options:
node_name: web_server
service_name: WebServer
service_name: web-server
# start webapp service
3:
action: "node_service_start"
action: "node-service-start"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
4:
action: "node_service_pause"
action: "node-service-pause"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
5:
action: "node_service_resume"
action: "node-service-resume"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
6:
action: "node_service_restart"
action: "node-service-restart"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
7:
action: "node_service_disable"
action: "node-service-disable"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
8:
action: "node_service_enable"
action: "node-service-enable"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
9: # check database.db file
action: "node_file_scan"
action: "node-file-scan"
options:
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_name: database_server
folder_name: database
file_name: database.db
11:
action: "node_file_delete"
action: "node-file-delete"
options:
node_name: database_server
folder_name: database
file_name: database.db
12:
action: "node_file_repair"
action: "node-file-repair"
options:
node_name: database_server
folder_name: database
file_name: database.db
13:
action: "node_service_fix"
action: "node-service-fix"
options:
node_name: database_server
service_name: DatabaseService
service_name: database-service
14:
action: "node_folder_scan"
action: "node-folder-scan"
options:
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_name: database_server
folder_name: database
16:
action: "node_folder_repair"
action: "node-folder-repair"
options:
node_name: database_server
folder_name: database
17:
action: "node_folder_restore"
action: "node-folder-restore"
options:
node_name: database_server
folder_name: database
18:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: domain_controller
19:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: domain_controller
20:
action: node_startup
action: node-startup
options:
node_name: domain_controller
21:
action: node_reset
action: node-reset
options:
node_name: domain_controller
22:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: web_server
23:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: web_server
24:
action: node_startup
action: node-startup
options:
node_name: web_server
25:
action: node_reset
action: node-reset
options:
node_name: web_server
26: # old action num: 18
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: database_server
27:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: database_server
28:
action: node_startup
action: node-startup
options:
node_name: database_server
29:
action: node_reset
action: node-reset
options:
node_name: database_server
30:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: backup_server
31:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: backup_server
32:
action: node_startup
action: node-startup
options:
node_name: backup_server
33:
action: node_reset
action: node-reset
options:
node_name: backup_server
34:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: security_suite
35:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: security_suite
36:
action: node_startup
action: node-startup
options:
node_name: security_suite
37:
action: node_reset
action: node-reset
options:
node_name: security_suite
38:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: client_1
39: # old action num: 19 # shutdown client 1
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: client_1
40: # old action num: 20
action: node_startup
action: node-startup
options:
node_name: client_1
41: # old action num: 21
action: node_reset
action: node-reset
options:
node_name: client_1
42:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: client_2
43:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: client_2
44:
action: node_startup
action: node-startup
options:
node_name: client_2
45:
action: node_reset
action: node-reset
options:
node_name: client_2
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "router_acl_add_rule"
46: # old action num: 22 # "acl: ADDRULE - Block outgoing traffic from client 1"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 1
@@ -410,8 +410,8 @@ agents:
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_add_rule"
47: # old action num: 23 # "acl: ADDRULE - Block outgoing traffic from client 2"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 2
@@ -424,7 +424,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 3
@@ -437,7 +437,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 4
@@ -450,7 +450,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
50: # old action num: 26
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 5
@@ -463,7 +463,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
51: # old action num: 27
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 6
@@ -476,132 +476,132 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
52: # old action num: 28
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 0
53: # old action num: 29
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 1
54: # old action num: 30
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 2
55: # old action num: 31
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 3
56: # old action num: 32
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 4
57: # old action num: 33
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 5
58: # old action num: 34
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 6
59: # old action num: 35
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 7
60: # old action num: 36
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 8
61: # old action num: 37
action: "router_acl_remove_rule"
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_name: domain_controller
nic_num: 1
63: # old action num: 39
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: domain_controller
nic_num: 1
64: # old action num: 40
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: web_server
nic_num: 1
65: # old action num: 41
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: web_server
nic_num: 1
66: # old action num: 42
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: database_server
nic_num: 1
67: # old action num: 43
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: database_server
nic_num: 1
68: # old action num: 44
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: backup_server
nic_num: 1
69: # old action num: 45
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: backup_server
nic_num: 1
70: # old action num: 46
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: security_suite
nic_num: 1
71: # old action num: 47
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: security_suite
nic_num: 1
72: # old action num: 48
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: security_suite
nic_num: 2
73: # old action num: 49
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: security_suite
nic_num: 2
74: # old action num: 50
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: client_1
nic_num: 1
75: # old action num: 51
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: client_1
nic_num: 1
76: # old action num: 52
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: client_2
nic_num: 1
77: # old action num: 53
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: client_2
nic_num: 1
@@ -611,19 +611,19 @@ agents:
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.40
options:
node_hostname: database_server
folder_name: database
file_name: database.db
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: client_1_green_user
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: client_2_green_user
@@ -693,7 +693,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
services:
- type: DNSServer
- type: dns-server
options:
domain_mapping:
arcd.com: 192.168.1.12 # web server
@@ -705,9 +705,9 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: WebServer
- type: web-server
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
@@ -719,10 +719,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: DatabaseService
- type: database-service
options:
backup_server_ip: 192.168.1.16
- type: FTPClient
- type: ftp-client
- hostname: backup_server
type: server
@@ -731,7 +731,7 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: FTPServer
- type: ftp-server
- hostname: security_suite
type: server
@@ -751,20 +751,20 @@ simulation:
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- type: DataManipulationBot
- type: data-manipulation-bot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- type: WebBrowser
- type: web-browser
options:
target_url: http://arcd.com/users/
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
services:
- type: DNSClient
- type: dns-client
- hostname: client_2
type: computer
@@ -773,20 +773,20 @@ simulation:
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- type: WebBrowser
- type: web-browser
options:
target_url: http://arcd.com/users/
- type: DataManipulationBot
- type: data-manipulation-bot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
services:
- type: DNSClient
- type: dns-client
links:
- endpoint_a_hostname: router_1

456
src/primaite/config/_package_data/data_manipulation_marl.yaml
View File

File diff suppressed because it is too large Load Diff

20
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml
View File

@@ -5,46 +5,46 @@ game:
agents:
- ref: RL_Agent
type: ProxyAgent
type: proxy-agent
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_shutdown
action: node-shutdown
options:
node_name: client_1
2:
action: node_shutdown
action: node-shutdown
options:
node_name: server
3:
action: node_startup
action: node-startup
options:
node_name: client_1
4:
action: node_startup
action: node-startup
options:
node_name: server
5:
action: host_nic_disable
action: host-nic-disable
options:
node_name: client_1
nic_num: 1
6:
action: host_nic_disable
action: host-nic-disable
options:
node_name: server
nic_num: 1
7:
action: host_nic_enable
action: host-nic-enable
options:
node_name: client_1
nic_num: 1
8:
action: host_nic_enable
action: host-nic-enable
options:
node_name: server
nic_num: 1

4
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_1.yaml
View File

@@ -1,5 +1,5 @@
server_services: &server_services
- type: DatabaseService
- type: database-service
client_applications: &client_applications
- type: DatabaseClient
- type: database-client

4
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_2.yaml
View File

@@ -1,5 +1,5 @@
server_services: &server_services
- type: FTPServer
- type: ftp-server
client_applications: &client_applications
- type: RansomwareScript
- type: ransomware-script

40
src/primaite/config/_package_data/multi_lan_internet_network_example.yaml
View File

@@ -20,10 +20,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -34,10 +34,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -103,7 +103,7 @@ simulation:
default_gateway: 8.8.8.1
services:
- ref: dns_server
type: DNSServer
type: dns-server
options:
domain_mapping:
sometech.ai: 94.10.180.6
@@ -150,7 +150,7 @@ simulation:
dst_ip: 94.10.180.6
dst_port: POSTGRES_SERVER
dst_wildcard_mask: 0.0.0.0
8: # Permit SomeTech DMZ to use ARP
8: # Permit SomeTech DMZ to use arp
action: PERMIT
src_port: ARP
dst_port: ARP
@@ -170,7 +170,7 @@ simulation:
dst_ip: 10.10.1.11
dst_port: POSTGRES_SERVER
dst_wildcard_mask: 0.0.0.0
8: # Permit SomeTech DMZ to use ARP
8: # Permit SomeTech DMZ to use arp
action: PERMIT
src_port: ARP
dst_port: ARP
@@ -197,9 +197,9 @@ simulation:
dns_server: 8.8.8.2
services:
- ref: web_server
type: WebServer
type: web-server
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
@@ -269,12 +269,12 @@ simulation:
action: PERMIT
src_port: HTTP
dst_port: HTTP
18: # Allow the SomeTech internal network to use ARP
18: # Allow the SomeTech internal network to use arp
action: PERMIT
src_ip: 10.10.0.0
src_wildcard_mask: 0.0.255.255
src_port: ARP
19: # Allow the SomeTech internal network to use ICMP
19: # Allow the SomeTech internal network to use icmp
action: PERMIT
src_ip: 10.10.0.0
src_wildcard_mask: 0.0.255.255
@@ -318,10 +318,10 @@ simulation:
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: DatabaseService
- type: database-service
options:
backup_server_ip: 10.10.1.12 # The some_tech_storage_srv server
- type: FTPClient
- type: ftp-client
- hostname: some_tech_storage_srv
type: server
@@ -330,7 +330,7 @@ simulation:
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: FTPServer
- type: ftp-server
- hostname: some_tech_hr_1
type: computer
@@ -339,10 +339,10 @@ simulation:
default_gateway: 10.10.3.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -353,10 +353,10 @@ simulation:
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -367,10 +367,10 @@ simulation:
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/

10
src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml
View File

@@ -1,7 +1,7 @@
agents: &greens
- ref: green_A
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.2
@@ -10,17 +10,17 @@ agents: &greens
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 1.0
options:
node_hostname: client

10
src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml
View File

@@ -1,7 +1,7 @@
agents: &greens
- ref: green_B
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.95
@@ -10,17 +10,17 @@ agents: &greens
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 1.0
options:
node_hostname: client

4
src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml
View File

@@ -1,11 +1,11 @@
reds: &reds
- ref: red_A
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client,]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 10
frequency: 10
variance: 0

4
src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml
View File

@@ -1,11 +1,11 @@
reds: &reds
- ref: red_B
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client_1]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 3
frequency: 2
variance: 1

34
src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml
View File

@@ -26,12 +26,12 @@ agents:
- ref: defender
team: BLUE
type: ProxyAgent
type: proxy-agent
observation_space:
type: CUSTOM
type: custom
options:
components:
- type: NODES
- type: nodes
label: NODES
options:
routers: []
@@ -46,7 +46,7 @@ agents:
include_num_access: false
include_nmne: true
- type: LINKS
- type: links
label: LINKS
options:
link_references:
@@ -56,48 +56,48 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_shutdown
action: node-shutdown
options:
node_name: client_1
2:
action: node_shutdown
action: node-shutdown
options:
node_name: server
3:
action: node_startup
action: node-startup
options:
node_name: client_1
4:
action: node_startup
action: node-startup
options:
node_name: server
5:
action: host_nic_disable
action: host-nic-disable
options:
node_name: client_1
nic_num: 1
6:
action: host_nic_disable
action: host-nic-disable
options:
node_name: server
nic_num: 1
7:
action: host_nic_enable
action: host-nic-enable
options:
node_name: client_1
nic_num: 1
8:
action: host_nic_enable
action: host-nic-enable
options:
node_name: server
nic_num: 1
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.40
options:
node_hostname: database_server
@@ -121,10 +121,10 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.3
- type: DataManipulationBot
- type: data-manipulation-bot
options:
server_ip: 192.168.1.3
payload: "DELETE"
@@ -139,7 +139,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
services:
- type: DatabaseService
- type: database-service
links:
- endpoint_a_hostname: client

19
src/primaite/game/agent/actions/abstract.py
View File

@@ -12,8 +12,6 @@ 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."""
@@ -22,13 +20,20 @@ class AbstractAction(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAction]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an action type.
:param discriminator: discriminator used to uniquely specify action types.
:type discriminator: str
:raises ValueError: When attempting to create an action with a name that is already in use.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {discriminator}")
cls._registry[discriminator] = cls
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:

28
src/primaite/game/agent/actions/acl.py
View File

@@ -21,9 +21,7 @@ __all__ = (
class ACLAddRuleAbstractAction(AbstractAction, ABC):
"""Base abstract class for ACL add rule actions."""
config: ConfigSchema = "ACLAddRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Configuration Schema base for ACL add rule abstract actions."""
src_ip: Union[IPV4Address, Literal["ALL"]]
@@ -37,19 +35,17 @@ class ACLAddRuleAbstractAction(AbstractAction, ABC):
dst_wildcard: Union[IPV4Address, Literal["NONE"]]
class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_abstract_action"):
"""Base abstract class for ACL remove rule actions."""
class ACLRemoveRuleAbstractAction(AbstractAction, ABC):
"""Base abstract class for acl remove rule actions."""
config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""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."""
class RouterACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="router-acl-add-rule"):
"""Action which adds a rule to a router's acl."""
config: "RouterACLAddRuleAction.ConfigSchema"
@@ -79,8 +75,8 @@ class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_ad
]
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_acl_remove_rule"):
"""Action which removes a rule from a router's ACL."""
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="router-acl-remove-rule"):
"""Action which removes a rule from a router's acl."""
config: "RouterACLRemoveRuleAction.ConfigSchema"
@@ -95,8 +91,8 @@ class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_
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."""
class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="firewall-acl-add-rule"):
"""Action which adds a rule to a firewall port's acl."""
config: "FirewallACLAddRuleAction.ConfigSchema"
@@ -130,8 +126,8 @@ class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_ac
]
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="firewall_acl_remove_rule"):
"""Action which removes a rule from a firewall port's ACL."""
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="firewall-acl-remove-rule"):
"""Action which removes a rule from a firewall port's acl."""
config: "FirewallACLRemoveRuleAction.ConfigSchema"

16
src/primaite/game/agent/actions/application.py
View File

@@ -23,9 +23,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
inherit from this base class.
"""
config: "NodeApplicationAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for Node Application actions."""
node_name: str
@@ -45,7 +43,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="node_application_execute"):
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, discriminator="node-application-execute"):
"""Action which executes an application."""
config: "NodeApplicationExecuteAction.ConfigSchema"
@@ -56,7 +54,7 @@ class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="no
verb: str = "execute"
class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_application_scan"):
class NodeApplicationScanAction(NodeApplicationAbstractAction, discriminator="node-application-scan"):
"""Action which scans an application."""
config: "NodeApplicationScanAction.ConfigSchema"
@@ -67,7 +65,7 @@ class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_
verb: str = "scan"
class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node_application_close"):
class NodeApplicationCloseAction(NodeApplicationAbstractAction, discriminator="node-application-close"):
"""Action which closes an application."""
config: "NodeApplicationCloseAction.ConfigSchema"
@@ -78,7 +76,7 @@ class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node
verb: str = "close"
class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_application_fix"):
class NodeApplicationFixAction(NodeApplicationAbstractAction, discriminator="node-application-fix"):
"""Action which fixes an application."""
config: "NodeApplicationFixAction.ConfigSchema"
@@ -89,7 +87,7 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_a
verb: str = "fix"
class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="node_application_install"):
class NodeApplicationInstallAction(NodeApplicationAbstractAction, discriminator="node-application-install"):
"""Action which installs an application."""
config: "NodeApplicationInstallAction.ConfigSchema"
@@ -113,7 +111,7 @@ class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="no
]
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, identifier="node_application_remove"):
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, discriminator="node-application-remove"):
"""Action which removes/uninstalls an application."""
config: "NodeApplicationRemoveAction.ConfigSchema"

28
src/primaite/game/agent/actions/file.py
View File

@@ -24,9 +24,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
only three parameters can inherit from this base class.
"""
config: "NodeFileAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Configuration Schema for NodeFileAbstractAction."""
node_name: str
@@ -38,7 +36,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
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 ["do-nothing"]
return [
"network",
"node",
@@ -52,7 +50,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
]
class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"):
class NodeFileCreateAction(NodeFileAbstractAction, discriminator="node-file-create"):
"""Action which creates a new file in a given folder."""
config: "NodeFileCreateAction.ConfigSchema"
@@ -67,7 +65,7 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
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 ["do-nothing"]
return [
"network",
"node",
@@ -81,7 +79,7 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
]
class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
class NodeFileScanAction(NodeFileAbstractAction, discriminator="node-file-scan"):
"""Action which scans a file."""
config: "NodeFileScanAction.ConfigSchema"
@@ -92,7 +90,7 @@ class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
verb: ClassVar[str] = "scan"
class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"):
class NodeFileDeleteAction(NodeFileAbstractAction, discriminator="node-file-delete"):
"""Action which deletes a file."""
config: "NodeFileDeleteAction.ConfigSchema"
@@ -106,7 +104,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
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 ["do-nothing"]
return [
"network",
"node",
@@ -119,7 +117,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
]
class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restore"):
class NodeFileRestoreAction(NodeFileAbstractAction, discriminator="node-file-restore"):
"""Action which restores a file."""
config: "NodeFileRestoreAction.ConfigSchema"
@@ -130,7 +128,7 @@ class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restor
verb: ClassVar[str] = "restore"
class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrupt"):
class NodeFileCorruptAction(NodeFileAbstractAction, discriminator="node-file-corrupt"):
"""Action which corrupts a file."""
config: "NodeFileCorruptAction.ConfigSchema"
@@ -141,7 +139,7 @@ class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrup
verb: ClassVar[str] = "corrupt"
class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"):
class NodeFileAccessAction(NodeFileAbstractAction, discriminator="node-file-access"):
"""Action which increases a file's access count."""
config: "NodeFileAccessAction.ConfigSchema"
@@ -155,7 +153,7 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
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 ["do-nothing"]
return [
"network",
"node",
@@ -167,7 +165,7 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
]
class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_checkhash"):
class NodeFileCheckhashAction(NodeFileAbstractAction, discriminator="node-file-checkhash"):
"""Action which checks the hash of a file."""
config: "NodeFileCheckhashAction.ConfigSchema"
@@ -178,7 +176,7 @@ class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_chec
verb: ClassVar[str] = "checkhash"
class NodeFileRepairAction(NodeFileAbstractAction, identifier="node_file_repair"):
class NodeFileRepairAction(NodeFileAbstractAction, discriminator="node-file-repair"):
"""Action which repairs a file."""
config: "NodeFileRepairAction.ConfigSchema"

18
src/primaite/game/agent/actions/folder.py
View File

@@ -22,9 +22,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
this base class.
"""
config: "NodeFolderAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NodeFolder actions."""
node_name: str
@@ -35,7 +33,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
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 ["do-nothing"]
return [
"network",
"node",
@@ -47,7 +45,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
]
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"):
class NodeFolderScanAction(NodeFolderAbstractAction, discriminator="node-folder-scan"):
"""Action which scans a folder."""
config: "NodeFolderScanAction.ConfigSchema"
@@ -58,7 +56,7 @@ class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_sca
verb: ClassVar[str] = "scan"
class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folder_checkhash"):
class NodeFolderCheckhashAction(NodeFolderAbstractAction, discriminator="node-folder-checkhash"):
"""Action which checks the hash of a folder."""
config: "NodeFolderCheckhashAction.ConfigSchema"
@@ -69,7 +67,7 @@ class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folde
verb: ClassVar[str] = "checkhash"
class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_repair"):
class NodeFolderRepairAction(NodeFolderAbstractAction, discriminator="node-folder-repair"):
"""Action which repairs a folder."""
config: "NodeFolderRepairAction.ConfigSchema"
@@ -80,7 +78,7 @@ class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_r
verb: ClassVar[str] = "repair"
class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_restore"):
class NodeFolderRestoreAction(NodeFolderAbstractAction, discriminator="node-folder-restore"):
"""Action which restores a folder."""
config: "NodeFolderRestoreAction.ConfigSchema"
@@ -91,7 +89,7 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_
verb: ClassVar[str] = "restore"
class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_create"):
class NodeFolderCreateAction(NodeFolderAbstractAction, discriminator="node-folder-create"):
"""Action which creates a new folder."""
config: "NodeFolderCreateAction.ConfigSchema"
@@ -105,7 +103,7 @@ class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_c
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 ["do-nothing"]
return [
"network",
"node",

10
src/primaite/game/agent/actions/host_nic.py
View File

@@ -16,9 +16,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
base class.
"""
config: "HostNICAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for HostNIC actions."""
node_name: str
@@ -29,7 +27,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
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 ["do-nothing"]
return [
"network",
"node",
@@ -40,7 +38,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
]
class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
class HostNICEnableAction(HostNICAbstractAction, discriminator="host-nic-enable"):
"""Action which enables a NIC."""
config: "HostNICEnableAction.ConfigSchema"
@@ -51,7 +49,7 @@ class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
verb: ClassVar[str] = "enable"
class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"):
class HostNICDisableAction(HostNICAbstractAction, discriminator="host-nic-disable"):
"""Action which disables a NIC."""
config: "HostNICDisableAction.ConfigSchema"

12
src/primaite/game/agent/actions/manager.py
View File

@@ -5,9 +5,9 @@ agents:
- name: agent_1
action_space:
actions:
- do_nothing
- node_service_start
- node_service_stop
- do-nothing
- node-service-start
- node-service-stop
action_map:
"""
@@ -24,18 +24,18 @@ from primaite.interface.request import RequestFormat
__all__ = ("DoNothingAction", "ActionManager")
class DoNothingAction(AbstractAction, identifier="do_nothing"):
class DoNothingAction(AbstractAction, discriminator="do-nothing"):
"""Do Nothing Action."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for do_nothingAction."""
type: str = "do_nothing"
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"]
return ["do-nothing"]
class _ActionMapItem(BaseModel):

13
src/primaite/game/agent/actions/network.py
View File

@@ -1,5 +1,6 @@
# © 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
@@ -8,12 +9,10 @@ from primaite.interface.request import RequestFormat
__all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction")
class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstract"):
class NetworkPortAbstractAction(AbstractAction, ABC):
"""Base class for Network port actions."""
config: "NetworkPortAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NetworkPort actions."""
target_nodename: str
@@ -24,7 +23,7 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac
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 ["do-nothing"]
return [
"network",
"node",
@@ -35,7 +34,7 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac
]
class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"):
class NetworkPortEnableAction(NetworkPortAbstractAction, discriminator="network-port-enable"):
"""Action which enables are port on a router or a firewall."""
config: "NetworkPortEnableAction.ConfigSchema"
@@ -46,7 +45,7 @@ class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_por
verb: ClassVar[str] = "enable"
class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"):
class NetworkPortDisableAction(NetworkPortAbstractAction, discriminator="network-port-disable"):
"""Action which disables are port on a router or a firewall."""
config: "NetworkPortDisableAction.ConfigSchema"

40
src/primaite/game/agent/actions/node.py
View File

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from abc import ABC, abstractmethod
from typing import ClassVar, List, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction
@@ -18,16 +18,14 @@ __all__ = (
)
class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
class NodeAbstractAction(AbstractAction, ABC):
"""
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):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for Node actions."""
node_name: str
@@ -39,7 +37,7 @@ class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
class NodeOSScanAction(NodeAbstractAction, discriminator="node-os-scan"):
"""Action which scans a node's OS."""
config: "NodeOSScanAction.ConfigSchema"
@@ -50,7 +48,7 @@ class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
verb: ClassVar[str] = "scan"
class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
class NodeShutdownAction(NodeAbstractAction, discriminator="node-shutdown"):
"""Action which shuts down a node."""
config: "NodeShutdownAction.ConfigSchema"
@@ -61,7 +59,7 @@ class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
verb: ClassVar[str] = "shutdown"
class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
class NodeStartupAction(NodeAbstractAction, discriminator="node-startup"):
"""Action which starts up a node."""
config: "NodeStartupAction.ConfigSchema"
@@ -72,7 +70,7 @@ class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
verb: ClassVar[str] = "startup"
class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
class NodeResetAction(NodeAbstractAction, discriminator="node-reset"):
"""Action which resets a node."""
config: "NodeResetAction.ConfigSchema"
@@ -83,12 +81,10 @@ class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
verb: ClassVar[str] = "reset"
class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_action"):
class NodeNMAPAbstractAction(AbstractAction, ABC):
"""Base class for NodeNMAP actions."""
config: "NodeNMAPAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration Schema for NodeNMAP actions."""
target_ip_address: Union[str, List[str]]
@@ -103,8 +99,8 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti
pass
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_scan"):
"""Action which performs an NMAP ping scan."""
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, discriminator="node-nmap-ping-scan"):
"""Action which performs an nmap ping scan."""
config: "NodeNMAPPingScanAction.ConfigSchema"
@@ -116,14 +112,14 @@ class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_
"node",
config.source_node,
"application",
"NMAP",
"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."""
class NodeNMAPPortScanAction(NodeNMAPAbstractAction, discriminator="node-nmap-port-scan"):
"""Action which performs an nmap port scan."""
config: "NodeNMAPPortScanAction.ConfigSchema"
@@ -143,7 +139,7 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
"node",
config.source_node,
"application",
"NMAP",
"nmap",
"port_scan",
{
"target_ip_address": config.target_ip_address,
@@ -154,8 +150,8 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
]
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_network_service_recon"):
"""Action which performs an NMAP network service recon (ping scan followed by port scan)."""
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, discriminator="node-network-service-recon"):
"""Action which performs an nmap network service recon (ping scan followed by port scan)."""
config: "NodeNetworkServiceReconAction.ConfigSchema"
@@ -174,7 +170,7 @@ class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_net
"node",
config.source_node,
"application",
"NMAP",
"nmap",
"network_service_recon",
{
"target_ip_address": config.target_ip_address,

25
src/primaite/game/agent/actions/service.py
View File

@@ -1,4 +1,5 @@
# © 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
@@ -17,15 +18,13 @@ __all__ = (
)
class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstract"):
class NodeServiceAbstractAction(AbstractAction, ABC):
"""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):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
node_name: str
service_name: str
verb: ClassVar[str]
@@ -36,7 +35,7 @@ class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstrac
return ["network", "node", config.node_name, "service", config.service_name, config.verb]
class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_scan"):
class NodeServiceScanAction(NodeServiceAbstractAction, discriminator="node-service-scan"):
"""Action which scans a service."""
config: "NodeServiceScanAction.ConfigSchema"
@@ -47,7 +46,7 @@ class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "scan"
class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_stop"):
class NodeServiceStopAction(NodeServiceAbstractAction, discriminator="node-service-stop"):
"""Action which stops a service."""
config: "NodeServiceStopAction.ConfigSchema"
@@ -58,7 +57,7 @@ class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "stop"
class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service_start"):
class NodeServiceStartAction(NodeServiceAbstractAction, discriminator="node-service-start"):
"""Action which starts a service."""
config: "NodeServiceStartAction.ConfigSchema"
@@ -69,7 +68,7 @@ class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "start"
class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service_pause"):
class NodeServicePauseAction(NodeServiceAbstractAction, discriminator="node-service-pause"):
"""Action which pauses a service."""
config: "NodeServicePauseAction.ConfigSchema"
@@ -80,7 +79,7 @@ class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "pause"
class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_service_resume"):
class NodeServiceResumeAction(NodeServiceAbstractAction, discriminator="node-service-resume"):
"""Action which resumes a service."""
config: "NodeServiceResumeAction.ConfigSchema"
@@ -91,7 +90,7 @@ class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "resume"
class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_service_restart"):
class NodeServiceRestartAction(NodeServiceAbstractAction, discriminator="node-service-restart"):
"""Action which restarts a service."""
config: "NodeServiceRestartAction.ConfigSchema"
@@ -102,7 +101,7 @@ class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "restart"
class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_service_disable"):
class NodeServiceDisableAction(NodeServiceAbstractAction, discriminator="node-service-disable"):
"""Action which disables a service."""
config: "NodeServiceDisableAction.ConfigSchema"
@@ -113,7 +112,7 @@ class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "disable"
class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_service_enable"):
class NodeServiceEnableAction(NodeServiceAbstractAction, discriminator="node-service-enable"):
"""Action which enables a service."""
config: "NodeServiceEnableAction.ConfigSchema"
@@ -124,7 +123,7 @@ class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "enable"
class NodeServiceFixAction(NodeServiceAbstractAction, identifier="node_service_fix"):
class NodeServiceFixAction(NodeServiceAbstractAction, discriminator="node-service-fix"):
"""Action which fixes a service."""
config: "NodeServiceFixAction.ConfigSchema"

26
src/primaite/game/agent/actions/session.py
View File

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from abc import ABC, abstractmethod
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
@@ -11,12 +11,10 @@ __all__ = (
)
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"):
class NodeSessionAbstractAction(AbstractAction, ABC):
"""Base class for NodeSession actions."""
config: "NodeSessionAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NodeSessionAbstractActions."""
node_name: str
@@ -33,7 +31,7 @@ class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstrac
pass
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_session_remote_login"):
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="node-session-remote-login"):
"""Action which performs a remote session login."""
config: "NodeSessionsRemoteLoginAction.ConfigSchema"
@@ -48,21 +46,21 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_
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 ["do-nothing"]
return [
"network",
"node",
config.node_name,
"service",
"Terminal",
"node_session_remote_login",
"terminal",
"node-session-remote-login",
config.username,
config.password,
config.remote_ip,
]
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node_session_remote_logoff"):
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, discriminator="node-session-remote-logoff"):
"""Action which performs a remote session logout."""
config: "NodeSessionsRemoteLogoutAction.ConfigSchema"
@@ -76,11 +74,11 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node
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]
return ["do-nothing"]
return ["network", "node", config.node_name, "service", "terminal", config.verb, config.remote_ip]
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"):
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, discriminator="node-account-change-password"):
"""Action which changes the password for a user."""
config: "NodeAccountChangePasswordAction.ConfigSchema"
@@ -100,7 +98,7 @@ class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="nod
"node",
config.node_name,
"service",
"UserManager",
"user-manager",
"change_password",
config.username,
config.current_password,

48
src/primaite/game/agent/actions/software.py
View File

@@ -22,7 +22,7 @@ __all__ = (
)
class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_ransomware_script"):
class ConfigureRansomwareScriptAction(AbstractAction, discriminator="configure-ransomware-script"):
"""Action which sets config parameters for a ransomware script on a node."""
config: "ConfigureRansomwareScriptAction.ConfigSchema"
@@ -39,16 +39,18 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_rans
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"]
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]
return ["network", "node", config.node_name, "application", "ransomware-script", "configure", data]
class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identifier="c2_server_ransomware_configure"):
class RansomwareConfigureC2ServerAction(
ConfigureRansomwareScriptAction, discriminator="c2-server-ransomware-configure"
):
"""Action which causes a C2 server to send a command to set options on a ransomware script remotely."""
@classmethod
@@ -56,10 +58,10 @@ class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identif
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]
return ["network", "node", config.node_name, "application", "c2-server", "ransomware_configure", data]
class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
class ConfigureDoSBotAction(AbstractAction, discriminator="configure-dos-bot"):
"""Action which sets config parameters for a DoS bot on a node."""
class ConfigSchema(AbstractAction.ConfigSchema):
@@ -88,10 +90,10 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
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]
return ["network", "node", config.node_name, "application", "dos-bot", "configure", data]
class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
class ConfigureC2BeaconAction(AbstractAction, discriminator="configure-c2-beacon"):
"""Action which configures a C2 Beacon based on the parameters given."""
class ConfigSchema(AbstractAction.ConfigSchema):
@@ -112,10 +114,10 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
masquerade_protocol=config.masquerade_protocol,
masquerade_port=config.masquerade_port,
)
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data]
return ["network", "node", config.node_name, "application", "c2-beacon", "configure", data]
class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"):
class NodeSendRemoteCommandAction(AbstractAction, discriminator="node-send-remote-command"):
"""Action which sends a terminal command to a remote node via SSH."""
config: "NodeSendRemoteCommandAction.ConfigSchema"
@@ -135,14 +137,14 @@ class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_c
"node",
config.node_name,
"service",
"Terminal",
"terminal",
"send_remote_command",
config.remote_ip,
{"command": config.command},
]
class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_command"):
class TerminalC2ServerAction(AbstractAction, discriminator="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"
@@ -160,7 +162,7 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
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"]
return ["do-nothing"]
command_model = {
"commands": config.commands,
@@ -168,10 +170,10 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model]
return ["network", "node", config.node_name, "application", "c2-server", "terminal_command", command_model]
class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ransomware_launch"):
class RansomwareLaunchC2ServerAction(AbstractAction, discriminator="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"
@@ -185,12 +187,12 @@ class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ranso
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"]
return ["do-nothing"]
# This action currently doesn't require any further configuration options.
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"]
return ["network", "node", config.node_name, "application", "c2-server", "ransomware_launch"]
class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfiltrate"):
class ExfiltrationC2ServerAction(AbstractAction, discriminator="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"
@@ -210,7 +212,7 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
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"]
return ["do-nothing"]
command_model = {
"target_file_name": config.target_file_name,
@@ -220,10 +222,10 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model]
return ["network", "node", config.node_name, "application", "c2-server", "exfiltrate", command_model]
class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"):
class ConfigureDatabaseClientAction(AbstractAction, discriminator="configure-database-client"):
"""Action which sets config parameters for a database client on a node."""
config: "ConfigureDatabaseClientAction.ConfigSchema"
@@ -239,6 +241,6 @@ class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_databa
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"]
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]
return ["network", "node", config.node_name, "application", "database-client", "configure", data]

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

@@ -68,7 +68,7 @@ class AbstractAgent(BaseModel, ABC):
)
reward_function: RewardFunction.ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema())
config: "AbstractAgent.ConfigSchema" = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
logger: AgentLog = AgentLog(agent_name="Abstract_Agent")
history: List[AgentHistoryItem] = []
@@ -79,13 +79,13 @@ class AbstractAgent(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {discriminator}")
cls._registry[discriminator] = cls
def model_post_init(self, __context: Any) -> None:
"""Overwrite the default empty action, observation, and rewards with ones defined through the config."""
@@ -130,7 +130,7 @@ class AbstractAgent(BaseModel, 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]) -> RequestFormat:
# this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator.
@@ -161,23 +161,23 @@ class AbstractAgent(BaseModel, ABC):
return agent_class(config=config)
class AbstractScriptedAgent(AbstractAgent, identifier="AbstractScriptedAgent"):
class AbstractScriptedAgent(AbstractAgent, ABC):
"""Base class for actors which generate their own behaviour."""
config: "AbstractScriptedAgent.ConfigSchema" = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema())
class ConfigSchema(AbstractAgent.ConfigSchema):
class ConfigSchema(AbstractAgent.ConfigSchema, ABC):
"""Configuration Schema for AbstractScriptedAgents."""
type: str = "AbstractScriptedAgent"
config: ConfigSchema = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema())
@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, identifier="ProxyAgent"):
class ProxyAgent(AbstractAgent, discriminator="proxy-agent"):
"""Agent that sends observations to an RL model and receives actions from that model."""
config: "ProxyAgent.ConfigSchema" = Field(default_factory=lambda: ProxyAgent.ConfigSchema())

2
src/primaite/game/agent/observations/acl_observation.py
View File

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class ACLObservation(AbstractObservation, identifier="ACL"):
class ACLObservation(AbstractObservation, discriminator="acl"):
"""ACL observation, provides information about access control lists within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/file_system_observations.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__)
class FileObservation(AbstractObservation, identifier="FILE"):
class FileObservation(AbstractObservation, discriminator="file"):
"""File observation, provides status information about a file within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -125,7 +125,7 @@ class FileObservation(AbstractObservation, identifier="FILE"):
)
class FolderObservation(AbstractObservation, identifier="FOLDER"):
class FolderObservation(AbstractObservation, discriminator="folder"):
"""Folder observation, provides status information about a folder within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/firewall_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
class FirewallObservation(AbstractObservation, discriminator="firewall"):
"""Firewall observation, provides status information about a firewall within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -181,7 +181,7 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
},
}
if self.include_users:
sess = firewall_state["services"]["UserSessionManager"]
sess = firewall_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

4
src/primaite/game/agent/observations/host_observations.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class HostObservation(AbstractObservation, identifier="HOST"):
class HostObservation(AbstractObservation, discriminator="host"):
"""Host observation, provides status information about a host within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -209,7 +209,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
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"]
sess = node_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

4
src/primaite/game/agent/observations/link_observation.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__)
class LinkObservation(AbstractObservation, identifier="LINK"):
class LinkObservation(AbstractObservation, discriminator="link"):
"""Link observation, providing information about a specific link within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -90,7 +90,7 @@ class LinkObservation(AbstractObservation, identifier="LINK"):
return cls(where=where)
class LinksObservation(AbstractObservation, identifier="LINKS"):
class LinksObservation(AbstractObservation, discriminator="links"):
"""Collection of link observations representing multiple links within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -12,7 +12,7 @@ from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.port import Port
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
class NICObservation(AbstractObservation, discriminator="network-interface"):
"""Status information about a network interface within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -227,7 +227,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
)
class PortObservation(AbstractObservation, identifier="PORT"):
class PortObservation(AbstractObservation, discriminator="port"):
"""Port observation, provides status information about a network port within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

2
src/primaite/game/agent/observations/node_observations.py
View File

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class NodesObservation(AbstractObservation, identifier="NODES"):
class NodesObservation(AbstractObservation, discriminator="nodes"):
"""Nodes observation, provides status information about nodes within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

18
src/primaite/game/agent/observations/observation_manager.py
View File

@@ -11,7 +11,7 @@ from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validat
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
class NestedObservation(AbstractObservation, identifier="CUSTOM"):
class NestedObservation(AbstractObservation, discriminator="custom"):
"""Observation type that allows combining other observations into a gymnasium.spaces.Dict space."""
class NestedObservationItem(BaseModel):
@@ -19,7 +19,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
model_config = ConfigDict(extra="forbid")
type: str
"""Select observation class. It maps to the identifier of the obs class by checking the registry."""
"""Select observation class. It maps to the discriminator of the obs class by checking the registry."""
label: str
"""Dict key in the final observation space."""
options: Dict
@@ -48,7 +48,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
def __init__(self, components: Dict[str, AbstractObservation]) -> None:
"""Initialise nested observation."""
self.components: Dict[str, AbstractObservation] = components
"""Maps label: observation object"""
"""Maps label observation object"""
self.default_observation = {label: obs.default_observation for label, obs in self.components.items()}
"""Default observation is just the default observations of constituents."""
@@ -84,7 +84,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
```yaml
observation_space:
- type: CUSTOM
- type: custom
options:
components:
@@ -119,7 +119,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
return cls(components=instances)
class NullObservation(AbstractObservation, identifier="NONE"):
class NullObservation(AbstractObservation, discriminator="none"):
"""Empty observation that acts as a placeholder."""
def __init__(self) -> None:
@@ -157,8 +157,8 @@ class ObservationManager(BaseModel):
"""Config Schema for Observation Manager."""
model_config = ConfigDict(extra="forbid")
type: str = "NONE"
"""Identifier name for the top-level observation."""
type: str = "none"
"""discriminator name for the top-level observation."""
options: AbstractObservation.ConfigSchema = Field(
default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True
)
@@ -187,7 +187,7 @@ class ObservationManager(BaseModel):
return data
# (TODO: duplicate default definition between here and the actual model)
obs_type = data["type"] if "type" in data else "NONE"
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
@@ -235,7 +235,7 @@ class ObservationManager(BaseModel):
:param config: Dictionary containing the configuration for this observation space.
If None, a blank observation space is created.
Otherwise, this must be a Dict with a type field and options field.
type: string that corresponds to one of the observation identifiers that are provided when subclassing
type: string that corresponds to one of the observation discriminators that are provided when subclassing
AbstractObservation
options: this must adhere to the chosen observation type's ConfigSchema nested class.
:type config: Dict

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

@@ -31,20 +31,20 @@ class AbstractObservation(ABC):
"""Initialise an observation. This method must be overwritten."""
self.default_observation: ObsType
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an observation type.
:param identifier: Identifier used to uniquely specify observation component types.
:type identifier: str
:param discriminator: discriminator used to uniquely specify observation component types.
:type discriminator: str
:raises ValueError: When attempting to create a component with a name that is already in use.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate observation component type {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate observation component type {discriminator}")
cls._registry[discriminator] = cls
@abstractmethod
def observe(self, state: Dict) -> Any:

4
src/primaite/game/agent/observations/router_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class RouterObservation(AbstractObservation, identifier="ROUTER"):
class RouterObservation(AbstractObservation, discriminator="router"):
"""Router observation, provides status information about a router within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -113,7 +113,7 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
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"]
sess = router_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

4
src/primaite/game/agent/observations/software_observation.py
View File

@@ -10,7 +10,7 @@ from primaite.game.agent.observations.observations import AbstractObservation, W
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
class ServiceObservation(AbstractObservation, discriminator="service"):
"""Service observation, shows status of a service in the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -73,7 +73,7 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"):
return cls(where=parent_where + ["services", config.service_name])
class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
class ApplicationObservation(AbstractObservation, discriminator="application"):
"""Application observation, shows the status of an application within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

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

@@ -12,7 +12,7 @@ the structure:
```yaml
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.5
options:
node_name: database_server
@@ -20,7 +20,7 @@ the structure:
file_name: database.db
- type: WEB_SERVER_404_PENALTY
- type: web-server-404-penalty
weight: 0.5
options:
node_name: web_server
@@ -43,25 +43,25 @@ _LOGGER = getLogger(__name__)
WhereType = Optional[Iterable[Union[str, int]]]
class AbstractReward(BaseModel):
class AbstractReward(BaseModel, ABC):
"""Base class for reward function components."""
config: "AbstractReward.ConfigSchema"
class ConfigSchema(BaseModel, ABC):
"""Config schema for AbstractReward."""
type: str = ""
config: ConfigSchema
_registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate reward {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate reward {discriminator}")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "AbstractReward":
@@ -92,7 +92,7 @@ class AbstractReward(BaseModel):
return 0.0
class DummyReward(AbstractReward, identifier="DUMMY"):
class DummyReward(AbstractReward, discriminator="dummy"):
"""Dummy reward function component which always returns 0.0."""
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
@@ -108,7 +108,7 @@ class DummyReward(AbstractReward, identifier="DUMMY"):
return 0.0
class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"):
class DatabaseFileIntegrity(AbstractReward, discriminator="database-file-integrity"):
"""Reward function component which rewards the agent for maintaining the integrity of a database file."""
config: "DatabaseFileIntegrity.ConfigSchema"
@@ -118,7 +118,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for DatabaseFileIntegrity."""
type: str = "DATABASE_FILE_INTEGRITY"
type: str = "database-file-integrity"
node_hostname: str
folder_name: str
file_name: str
@@ -161,7 +161,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
return 0
class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
class WebServer404Penalty(AbstractReward, discriminator="web-server-404-penalty"):
"""Reward function component which penalises the agent when the web server returns a 404 error."""
config: "WebServer404Penalty.ConfigSchema"
@@ -171,7 +171,7 @@ class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for WebServer404Penalty."""
type: str = "WEB_SERVER_404_PENALTY"
type: str = "web-server-404-penalty"
node_hostname: str
service_name: str
sticky: bool = True
@@ -215,7 +215,7 @@ class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
return self.reward
class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_PENALTY"):
class WebpageUnavailablePenalty(AbstractReward, discriminator="webpage-unavailable-penalty"):
"""Penalises the agent when the web browser fails to fetch a webpage."""
config: "WebpageUnavailablePenalty.ConfigSchema"
@@ -225,7 +225,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for WebpageUnavailablePenalty."""
type: str = "WEBPAGE_UNAVAILABLE_PENALTY"
type: str = "webpage-unavailable-penalty"
node_hostname: str = ""
sticky: bool = True
@@ -248,7 +248,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
"nodes",
self.config.node_hostname,
"applications",
"WebBrowser",
"web-browser",
]
web_browser_state = access_from_nested_dict(state, self.location_in_state)
@@ -261,7 +261,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
"node",
self.config.node_hostname,
"application",
"WebBrowser",
"web-browser",
"execute",
]
@@ -289,7 +289,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
return self.reward
class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"):
class GreenAdminDatabaseUnreachablePenalty(AbstractReward, discriminator="green-admin-database-unreachable-penalty"):
"""Penalises the agent when the green db clients fail to connect to the database."""
config: "GreenAdminDatabaseUnreachablePenalty.ConfigSchema"
@@ -298,7 +298,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for GreenAdminDatabaseUnreachablePenalty."""
type: str = "GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"
type: str = "green-admin-database-unreachable-penalty"
node_hostname: str
sticky: bool = True
@@ -322,7 +322,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
"node",
self.config.node_hostname,
"application",
"DatabaseClient",
"database-client",
"execute",
]
@@ -339,7 +339,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
return self.reward
class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
class SharedReward(AbstractReward, discriminator="shared-reward"):
"""Adds another agent's reward to the overall reward."""
config: "SharedReward.ConfigSchema"
@@ -347,7 +347,7 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
class ConfigSchema(AbstractReward.ConfigSchema):
"""Config schema for SharedReward."""
type: str = "SHARED_REWARD"
type: str = "shared-reward"
agent_name: str
def default_callback(agent_name: str) -> Never:
@@ -376,17 +376,17 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
return self.callback(self.config.agent_name)
class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"):
"""Apply a negative reward when taking any action except do_nothing."""
class ActionPenalty(AbstractReward, discriminator="action-penalty"):
"""Apply a negative reward when taking any action except do-nothing."""
config: "ActionPenalty.ConfigSchema"
class ConfigSchema(AbstractReward.ConfigSchema):
"""Config schema for ActionPenalty.
:param action_penalty: Reward to give agents for taking any action except do_nothing
: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
:param do_nothing_penalty: Reward to give agent for taking the do-nothing action
:type do_nothing_penalty: float
"""
@@ -403,7 +403,7 @@ class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"):
:return: Reward value
:rtype: float
"""
if last_action_response.action == "do_nothing":
if last_action_response.action == "do-nothing":
return self.config.do_nothing_penalty
else:

10
src/primaite/game/agent/scripted_agents/abstract_tap.py
View File

@@ -2,7 +2,7 @@
from __future__ import annotations
import random
from abc import abstractmethod
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple
from gymnasium.core import ObsType
@@ -13,21 +13,21 @@ from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "AbstractTAPAgent"
class AbstractTAPAgent(PeriodicAgent, identifier="AbstractTAP"):
class AbstractTAPAgent(PeriodicAgent, ABC):
"""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):
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema, ABC):
"""Schema for the `agent_settings` part of the agent config."""
possible_starting_nodes: List[str] = Field(default_factory=list)
class ConfigSchema(PeriodicAgent.ConfigSchema):
class ConfigSchema(PeriodicAgent.ConfigSchema, ABC):
"""Configuration schema for Abstract TAP agents."""
type: str = "AbstractTAP"
type: str = "abstract-tap"
agent_settings: AbstractTAPAgent.AgentSettingsSchema = Field(
default_factory=lambda: AbstractTAPAgent.AgentSettingsSchema()
)

12
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -6,21 +6,19 @@ from pydantic import Field
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "DataManipulationAgent"
class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"):
class DataManipulationAgent(PeriodicAgent, discriminator="red-database-corrupting-agent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
target_application: str = "DataManipulationBot"
target_application: str = "data-manipulation-bot"
class ConfigSchema(PeriodicAgent.ConfigSchema):
"""Configuration Schema for DataManipulationAgent."""
type: str = "RedDatabaseCorruptingAgent"
type: str = "red-database-corrupting-agent"
agent_settings: "DataManipulationAgent.AgentSettingsSchema" = Field(
default_factory=lambda: DataManipulationAgent.AgentSettingsSchema()
)
@@ -43,13 +41,13 @@ class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgen
"""
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
return "do-nothing", {}
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", {
return "node-application-execute", {
"node_name": self.start_node,
"application_name": self.config.agent_settings.target_application,
}

4
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = "ProbabilisticAgent"
class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"):
class ProbabilisticAgent(AbstractScriptedAgent, discriminator="probabilistic-agent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535)))
@@ -46,7 +46,7 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration schema for Probabilistic Agent."""
type: str = "ProbabilisticAgent"
type: str = "probabilistic-agent"
agent_settings: "ProbabilisticAgent.AgentSettingsSchema" = Field(
default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema()
)

12
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -11,7 +11,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = ("RandomAgent", "PeriodicAgent")
class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
class RandomAgent(AbstractScriptedAgent, discriminator="random-agent"):
"""Agent that ignores its observation and acts completely at random."""
config: "RandomAgent.ConfigSchema" = Field(default_factory=lambda: RandomAgent.ConfigSchema())
@@ -19,7 +19,7 @@ class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Random Agents."""
type: str = "RandomAgent"
type: str = "random-agent"
def get_action(self) -> Tuple[str, Dict]:
"""Sample the action space randomly.
@@ -34,7 +34,7 @@ class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
return self.action_manager.get_action(self.action_manager.space.sample())
class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
class PeriodicAgent(AbstractScriptedAgent, discriminator="periodic-agent"):
"""Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""
config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema())
@@ -72,7 +72,7 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Periodic Agent."""
type: str = "PeriodicAgent"
type: str = "periodic-agent"
"""Name of the agent."""
agent_settings: "PeriodicAgent.AgentSettingsSchema" = Field(
default_factory=lambda: PeriodicAgent.AgentSettingsSchema()
@@ -113,9 +113,9 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
self._set_next_execution_timestep(
timestep + self.config.agent_settings.frequency, self.config.agent_settings.variance
)
return "node_application_execute", {
return "node-application-execute", {
"node_name": self.start_node,
"application_name": self.config.agent_settings.target_application,
}
return "do_nothing", {}
return "do-nothing", {}

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

@@ -268,7 +268,7 @@ class PrimaiteGame:
new_node = None
if n_type in Node._registry:
if n_type == "wireless_router":
if n_type == "wireless-router":
node_cfg["airspace"] = net.airspace
new_node = Node._registry[n_type].from_config(config=node_cfg)
else:
@@ -288,8 +288,8 @@ class PrimaiteGame:
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
if "users" in node_cfg and new_node.software_manager.software.get("user-manager"):
user_manager: UserManager = new_node.software_manager.software["user-manager"] # noqa
for user_cfg in node_cfg["users"]:
user_manager.add_user(**user_cfg, bypass_can_perform_action=True)
@@ -321,7 +321,7 @@ 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, software_config=service_cfg.get("options", {}))
new_service = new_node.software_manager.software[service_class.__name__]
new_service = new_node.software_manager.software[service_type]
# fixing duration for the service
if "fixing_duration" in service_cfg.get("options", {}):

6
src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb
View File

@@ -109,7 +109,7 @@
" - 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",
@@ -416,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",
@@ -435,7 +435,7 @@
" ...\n",
" action_map:\n",
" 5:\n",
" action: c2_server_ransomware_configure\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_id: 1\n",
" config:\n",

2
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -172,7 +172,7 @@
"\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application, identifier=\"MSPaint\"):\n",
"class MSPaint(Application, discriminator=\"MSPaint\"):\n",
" def describe_state(self):\n",
" return super().describe_state()"
]

2
src/primaite/simulator/core.py
View File

@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
..code::python
class WebBrowser(Application, identifier="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager() # all requests generic to any Application get initialised
rm.add_request(...) # initialise any requests specific to the web browser

4
src/primaite/simulator/file_system/file.py
View File

@@ -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:

4
src/primaite/simulator/file_system/folder.py
View File

@@ -396,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:

2
src/primaite/simulator/network/container.py
View File

@@ -180,7 +180,7 @@ class Network(SimComponent):
table.align = "l"
table.title = "Nodes"
for node in self.nodes.values():
table.add_row((node.config.hostname, type(node)._identifier, node.operating_state.name))
table.add_row((node.config.hostname, type(node)._discriminator, node.operating_state.name))
print(table)
if ip_addresses:

24
src/primaite/simulator/network/creation.py
View File

@@ -22,7 +22,7 @@ class NetworkNodeAdder(BaseModel):
Here is a template that users can use to define custom node adders:
```
class YourNodeAdder(NetworkNodeAdder, identifier="your_name"):
class YourNodeAdder(NetworkNodeAdder, discriminator="your-name"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
property_1 : str
property_2 : int
@@ -40,8 +40,8 @@ class NetworkNodeAdder(BaseModel):
"""
Base schema for node adders.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The identifier is used
by the from_config method to select the correct node adder at runtime.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The discriminator is
used by the from_config method to select the correct node adder at runtime.
"""
model_config = ConfigDict(extra="forbid")
@@ -50,20 +50,20 @@ class NetworkNodeAdder(BaseModel):
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str], **kwargs: Any) -> None:
"""
Register a network node adder class.
:param identifier: Unique name for the node adder to use for matching against primaite config entries.
:type identifier: str
:param discriminator: Unique name for the node adder to use for matching against primaite config entries.
:type discriminator: str
:raises ValueError: When attempting to register a name that is already reserved.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate node adder {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate node adder {discriminator}")
cls._registry[discriminator] = cls
@classmethod
@abstractmethod
@@ -99,13 +99,13 @@ class NetworkNodeAdder(BaseModel):
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"):
class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
"""Creates an office LAN."""
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
"""Configuration schema for OfficeLANAdder."""
type: Literal["office_lan"] = "office_lan"
type: Literal["office-lan"] = "office-lan"
lan_name: str
"""Name of lan used for generating hostnames for new nodes."""
subnet_base: int

47
src/primaite/simulator/network/hardware/base.py
View File

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump()
class UserManager(Service, identifier="UserManager"):
class UserManager(Service, discriminator="user-manager"):
"""
Manages users within the PrimAITE system, handling creation, authentication, and administration.
@@ -836,7 +836,7 @@ class UserManager(Service, identifier="UserManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserManager."""
type: str = "UserManager"
type: str = "user-manager"
config: "UserManager.ConfigSchema" = Field(default_factory=lambda: UserManager.ConfigSchema())
@@ -849,7 +849,7 @@ class UserManager(Service, identifier="UserManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserManager"
kwargs["name"] = "user-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1037,7 +1037,7 @@ class UserManager(Service, identifier="UserManager"):
@property
def _user_session_manager(self) -> "UserSessionManager":
return self.software_manager.software["UserSessionManager"] # noqa
return self.software_manager.software["user-session-manager"] # noqa
class UserSession(SimComponent):
@@ -1137,7 +1137,7 @@ class RemoteUserSession(UserSession):
return state
class UserSessionManager(Service, identifier="UserSessionManager"):
class UserSessionManager(Service, discriminator="user-session-manager"):
"""
Manages user sessions on a Node, including local and remote sessions.
@@ -1147,7 +1147,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserSessionManager."""
type: str = "UserSessionManager"
type: str = "user-session-manager"
config: "UserSessionManager.ConfigSchema" = Field(default_factory=lambda: UserSessionManager.ConfigSchema())
@@ -1179,7 +1179,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserSessionManager"
kwargs["name"] = "user-session-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1289,7 +1289,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
:return: The UserManager instance.
"""
return self.software_manager.software["UserManager"] # noqa
return self.software_manager.software["user-manager"] # noqa
def pre_timestep(self, timestep: int) -> None:
"""Apply any pre-timestep logic that helps make sure we have the correct observations."""
@@ -1522,8 +1522,9 @@ class Node(SimComponent, ABC):
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
_identifier: ClassVar[str] = "unknown"
"""Identifier for this particular class, used for printing and logging. Each subclass redefines this."""
# TODO: this should not be set for abstract classes.
_discriminator: ClassVar[str]
"""discriminator for this particular class, used for printing and logging. Each subclass redefines this."""
class ConfigSchema(BaseModel, ABC):
"""Configuration Schema for Node based classes."""
@@ -1586,22 +1587,22 @@ class Node(SimComponent, ABC):
obj = cls(config=cls.ConfigSchema(**config))
return obj
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a node type.
:param identifier: Uniquely specifies an node class by name. Used for finding items by config.
:type identifier: str
:param discriminator: Uniquely specifies an node class by name. Used for finding items by config.
:type discriminator: str
:raises ValueError: When attempting to register an node with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
identifier = identifier.lower()
if identifier in cls._registry:
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
cls._identifier = identifier
discriminator = discriminator.lower()
if discriminator in cls._registry:
raise ValueError(f"Tried to define new node {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
cls._discriminator = discriminator
def __init__(self, **kwargs):
"""
@@ -1637,17 +1638,17 @@ class Node(SimComponent, ABC):
@property
def user_manager(self) -> Optional[UserManager]:
"""The Nodes User Manager."""
return self.software_manager.software.get("UserManager") # noqa
return self.software_manager.software.get("user-manager") # noqa
@property
def user_session_manager(self) -> Optional[UserSessionManager]:
"""The Nodes User Session Manager."""
return self.software_manager.software.get("UserSessionManager") # noqa
return self.software_manager.software.get("user-session-manager") # noqa
@property
def terminal(self) -> Optional[Terminal]:
"""The Nodes Terminal."""
return self.software_manager.software.get("Terminal")
"""The Node's Terminal."""
return self.software_manager.software.get("terminal")
def local_login(self, username: str, password: str) -> Optional[str]:
"""

4
src/primaite/simulator/network/hardware/nodes/host/computer.py
View File

@@ -7,7 +7,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(HostNode, identifier="computer"):
class Computer(HostNode, discriminator="computer"):
"""
A basic Computer class.
@@ -35,7 +35,7 @@ class Computer(HostNode, identifier="computer"):
* Web Browser
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient}
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "ftp-client": FTPClient}
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Computer class."""

26
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -267,7 +267,7 @@ class NIC(IPWiredNetworkInterface):
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
class HostNode(Node, identifier="HostNode"):
class HostNode(Node, discriminator="host-node"):
"""
Represents a host node in the network.
@@ -314,14 +314,14 @@ class HostNode(Node, identifier="HostNode"):
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"HostARP": HostARP,
"ICMP": ICMP,
"DNSClient": DNSClient,
"NTPClient": NTPClient,
"WebBrowser": WebBrowser,
"NMAP": NMAP,
"UserSessionManager": UserSessionManager,
"UserManager": UserManager,
"Terminal": Terminal,
"icmp": ICMP,
"dns-client": DNSClient,
"ntp-client": NTPClient,
"web-browser": WebBrowser,
"nmap": NMAP,
"user-session-manager": UserSessionManager,
"user-manager": UserManager,
"terminal": Terminal,
}
"""List of system software that is automatically installed on nodes."""
@@ -351,7 +351,7 @@ class HostNode(Node, identifier="HostNode"):
:return: NMAP application installed on the Node.
:rtype: Optional[NMAP]
"""
return self.software_manager.software.get("NMAP")
return self.software_manager.software.get("nmap")
@property
def arp(self) -> Optional[ARP]:
@@ -361,7 +361,7 @@ class HostNode(Node, identifier="HostNode"):
:return: ARP Cache for given HostNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")
def default_gateway_hello(self):
"""
@@ -393,8 +393,8 @@ class HostNode(Node, identifier="HostNode"):
dst_port = frame.udp.dst_port
can_accept_nmap = False
if self.software_manager.software.get("NMAP"):
if self.software_manager.software["NMAP"].operating_state == ApplicationOperatingState.RUNNING:
if self.software_manager.software.get("nmap"):
if self.software_manager.software["nmap"].operating_state == ApplicationOperatingState.RUNNING:
can_accept_nmap = True
accept_nmap = can_accept_nmap and frame.payload.__class__.__name__ == "PortScanPayload"

4
src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -5,7 +5,7 @@ from pydantic import Field
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
class Server(HostNode, identifier="server"):
class Server(HostNode, discriminator="server"):
"""
A basic Server class.
@@ -41,7 +41,7 @@ class Server(HostNode, identifier="server"):
config: ConfigSchema = Field(default_factory=lambda: Server.ConfigSchema())
class Printer(HostNode, identifier="printer"):
class Printer(HostNode, discriminator="printer"):
"""Printer? I don't even know her!."""
# TODO: Implement printer-specific behaviour

2
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -26,7 +26,7 @@ DMZ_PORT_ID: Final[int] = 3
"""The Firewall port ID of the DMZ port."""
class Firewall(Router, identifier="firewall"):
class Firewall(Router, discriminator="firewall"):
"""
A Firewall class that extends the functionality of a Router.

4
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -7,7 +7,7 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node, identifier="NetworkNode"):
class NetworkNode(Node, discriminator="network-node"):
"""
Represents an abstract base class for a network node that can receive and process network frames.
@@ -40,4 +40,4 @@ class NetworkNode(Node, identifier="NetworkNode"):
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")

2
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
class Router(NetworkNode, identifier="router"):
class Router(NetworkNode, discriminator="router"):
"""
Represents a network router, managing routing and forwarding of IP packets across network interfaces.

2
src/primaite/simulator/network/hardware/nodes/network/switch.py
View File

@@ -88,7 +88,7 @@ class SwitchPort(WiredNetworkInterface):
return False
class Switch(NetworkNode, identifier="switch"):
class Switch(NetworkNode, discriminator="switch"):
"""A class representing a Layer 2 network switch."""
network_interfaces: Dict[str, SwitchPort] = {}

2
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
)
class WirelessRouter(Router, identifier="wireless_router"):
class WirelessRouter(Router, discriminator="wireless-router"):
"""
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.

16
src/primaite/simulator/network/networks.py
View File

@@ -170,13 +170,13 @@ def arcd_uc2_network() -> Network:
client_1.power_on()
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
db_client_1: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
db_client_1: DatabaseClient = client_1.software_manager.software.get("database-client")
db_client_1.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_1.run()
web_browser_1 = client_1.software_manager.software.get("WebBrowser")
web_browser_1 = client_1.software_manager.software.get("web-browser")
web_browser_1.target_url = "http://arcd.com/users/"
client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("data-manipulation-bot")
db_manipulation_bot.configure(
server_ip_address=IPv4Address("192.168.1.14"),
payload="DELETE",
@@ -199,10 +199,10 @@ def arcd_uc2_network() -> Network:
client_2.power_on()
client_2.software_manager.install(DatabaseClient)
db_client_2 = client_2.software_manager.software.get("DatabaseClient")
db_client_2 = client_2.software_manager.software.get("database-client")
db_client_2.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_2.run()
web_browser_2 = client_2.software_manager.software.get("WebBrowser")
web_browser_2 = client_2.software_manager.software.get("web-browser")
web_browser_2.target_url = "http://arcd.com/users/"
network.connect(
endpoint_b=client_2.network_interface[1],
@@ -244,7 +244,7 @@ def arcd_uc2_network() -> Network:
network.connect(endpoint_b=database_server.network_interface[1], endpoint_a=switch_1.network_interface[3])
database_server.software_manager.install(DatabaseService)
database_service: DatabaseService = database_server.software_manager.software.get("DatabaseService") # noqa
database_service: DatabaseService = database_server.software_manager.software.get("database-service") # noqa
database_service.start()
database_service.configure_backup(backup_server=IPv4Address("192.168.1.16"))
@@ -264,7 +264,7 @@ def arcd_uc2_network() -> Network:
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
database_client: DatabaseClient = web_server.software_manager.software.get("database-client")
database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
network.connect(endpoint_b=web_server.network_interface[1], endpoint_a=switch_1.network_interface[2])
database_client.run()
@@ -273,7 +273,7 @@ def arcd_uc2_network() -> Network:
web_server.software_manager.install(WebServer)
# register the web_server to a domain
dns_server_service: DNSServer = domain_controller.software_manager.software.get("DNSServer") # noqa
dns_server_service: DNSServer = domain_controller.software_manager.software.get("dns-server") # noqa
dns_server_service.dns_register("arcd.com", web_server.network_interface[1].ip_address)
# Backup Server

4
src/primaite/simulator/sim_container.py
View File

@@ -38,8 +38,8 @@ class Simulation(SimComponent):
rm.add_request("network", RequestType(func=self.network._request_manager))
# pass through domain requests to the domain object
rm.add_request("domain", RequestType(func=self.domain._request_manager))
# if 'do_nothing' is requested, just return a success
rm.add_request("do_nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
# if 'do-nothing' is requested, just return a success
rm.add_request("do-nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
return rm
def describe_state(self) -> Dict:

14
src/primaite/simulator/system/applications/application.py
View File

@@ -53,20 +53,20 @@ class Application(IOSoftware, ABC):
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: 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: Optional[str]
:param discriminator: Uniquely specifies an application class by name. Used for finding items by config.
:type discriminator: Optional[str]
:raises ValueError: When attempting to register an application with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator 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
if discriminator in cls._registry:
raise ValueError(f"Tried to define new application {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "Application":

8
src/primaite/simulator/system/applications/database_client.py
View File

@@ -37,7 +37,7 @@ class DatabaseClientConnection(BaseModel):
@property
def client(self) -> Optional[DatabaseClient]:
"""The DatabaseClient that holds this connection."""
return self.parent_node.software_manager.software.get("DatabaseClient")
return self.parent_node.software_manager.software.get("database-client")
def query(self, sql: str) -> bool:
"""
@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
return str(self)
class DatabaseClient(Application, identifier="DatabaseClient"):
class DatabaseClient(Application, discriminator="database-client"):
"""
A DatabaseClient application.
@@ -72,7 +72,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for DatabaseClient."""
type: str = "DatabaseClient"
type: str = "database-client"
db_server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
@@ -97,7 +97,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
"""Native Client Connection for using the client directly (similar to psql in a terminal)."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
kwargs["name"] = "database-client"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

6
src/primaite/simulator/system/applications/nmap.py
View File

@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
return state
class NMAP(Application, identifier="NMAP"):
class NMAP(Application, discriminator="nmap"):
"""
A class representing the NMAP application for network scanning.
@@ -55,7 +55,7 @@ class NMAP(Application, identifier="NMAP"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for NMAP."""
type: str = "NMAP"
type: str = "nmap"
config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema())
@@ -70,7 +70,7 @@ class NMAP(Application, identifier="NMAP"):
}
def __init__(self, **kwargs):
kwargs["name"] = "NMAP"
kwargs["name"] = "nmap"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)

8
src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py
View File

@@ -162,11 +162,11 @@ class AbstractC2(Application):
:return: An FTPClient object is successful, else None
:rtype: union[FTPClient, None]
"""
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("FTPClient")
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("ftp-client")
if ftp_client is None:
self.sys_log.warning(f"{self.__class__.__name__}: No FTPClient. Attempting to install.")
self.software_manager.install(FTPClient)
ftp_client = self.software_manager.software.get("FTPClient")
ftp_client = self.software_manager.software.get("ftp-client")
# Force start if the service is stopped.
if ftp_client.operating_state == ServiceOperatingState.STOPPED:
@@ -189,11 +189,11 @@ class AbstractC2(Application):
:return: An FTPServer object is successful, else None
:rtype: Optional[FTPServer]
"""
ftp_server: Optional[FTPServer] = self.software_manager.software.get("FTPServer")
ftp_server: Optional[FTPServer] = self.software_manager.software.get("ftp-server")
if ftp_server is None:
self.sys_log.warning(f"{self.__class__.__name__}:No FTPServer installed. Attempting to install FTPServer.")
self.software_manager.install(FTPServer)
ftp_server = self.software_manager.software.get("FTPServer")
ftp_server = self.software_manager.software.get("ftp-server")
# Force start if the service is stopped.
if ftp_server.operating_state == ServiceOperatingState.STOPPED:

26
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -17,7 +17,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
class C2Beacon(AbstractC2, identifier="C2Beacon"):
class C2Beacon(AbstractC2, discriminator="c2-beacon"):
"""
C2 Beacon Application.
@@ -39,7 +39,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Beacon."""
type: str = "C2Beacon"
type: str = "c2-beacon"
c2_server_ip_address: Optional[IPV4Address] = None
keep_alive_frequency: int = 5
masquerade_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"]
@@ -54,13 +54,13 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
"The currently in use terminal session."
def __init__(self, **kwargs):
kwargs["name"] = "C2Beacon"
kwargs["name"] = "c2-beacon"
super().__init__(**kwargs)
@property
def _host_terminal(self) -> Optional[Terminal]:
"""Return the Terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("Terminal")
"""Return the terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("terminal")
if host_terminal is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a terminal on its host.")
return host_terminal
@@ -68,7 +68,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
@property
def _host_ransomware_script(self) -> RansomwareScript:
"""Return the RansomwareScript that is installed on the same machine as the C2 Beacon."""
ransomware_script: RansomwareScript = self.software_manager.software.get("RansomwareScript")
ransomware_script: RansomwareScript = self.software_manager.software.get("ransomware-script")
if ransomware_script is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find installed ransomware on its host.")
return ransomware_script
@@ -300,7 +300,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = RansomwareOpts.model_validate(payload.payload)
@@ -324,7 +324,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ransomware_script is None:
@@ -351,7 +351,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns a tuple containing Request Response returned by the Terminal execute method.
:return: Returns a tuple containing Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ftp_server is None:
@@ -372,7 +372,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
)
# Using the terminal to start the FTP Client on the remote machine.
self.terminal_session.execute(command=["service", "start", "FTPClient"])
self.terminal_session.execute(command=["service", "start", "ftp-client"])
# Need to supply to the FTP Client the C2 Beacon's host IP.
host_network_interfaces = self.software_manager.node.network_interfaces
@@ -430,7 +430,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
# Using the terminal to send the target data back to the C2 Beacon.
exfil_response: RequestResponse = RequestResponse.from_bool(
self.terminal_session.execute(command=["service", "FTPClient", "send", ftp_opts])
self.terminal_session.execute(command=["service", "ftp-client", "send", ftp_opts])
)
# Validating that we successfully received the target data.
@@ -472,14 +472,14 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
def _command_terminal(self, payload: C2Packet) -> RequestResponse:
"""
C2 Command: Terminal.
C2 Command: terminal.
Creates a request that executes a terminal command.
This request is then sent to the terminal service in order to be executed.
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = TerminalOpts.model_validate(payload.payload)

6
src/primaite/simulator/system/applications/red_applications/c2/c2_server.py
View File

@@ -16,7 +16,7 @@ from primaite.simulator.system.applications.red_applications.c2 import (
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
class C2Server(AbstractC2, identifier="C2Server"):
class C2Server(AbstractC2, discriminator="c2-server"):
"""
C2 Server Application.
@@ -37,7 +37,7 @@ class C2Server(AbstractC2, identifier="C2Server"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Server."""
type: str = "C2Server"
type: str = "c2-server"
config: ConfigSchema = Field(default_factory=lambda: C2Server.ConfigSchema())
@@ -125,7 +125,7 @@ class C2Server(AbstractC2, identifier="C2Server"):
return rm
def __init__(self, **kwargs):
kwargs["name"] = "C2Server"
kwargs["name"] = "c2-server"
super().__init__(**kwargs)
self.run()

8
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -40,13 +40,13 @@ class DataManipulationAttackStage(IntEnum):
"Signifies that the attack has failed."
class DataManipulationBot(Application, identifier="DataManipulationBot"):
class DataManipulationBot(Application, discriminator="data-manipulation-bot"):
"""A bot that simulates a script which performs a SQL injection attack."""
class ConfigSchema(Application.ConfigSchema):
"""Configuration schema for DataManipulationBot."""
type: str = "DataManipulationBot"
type: str = "data-manipulation-bot"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "DELETE"
@@ -64,7 +64,7 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
"Whether to repeat attacking once finished."
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulationBot"
kwargs["name"] = "data-manipulation-bot"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -92,7 +92,7 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the DataManipulationBot."""
db_client = self.software_manager.software.get("DatabaseClient")
db_client = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

6
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -32,13 +32,13 @@ class DoSAttackStage(IntEnum):
"Attack is completed."
class DoSBot(DatabaseClient, identifier="DoSBot"):
class DoSBot(DatabaseClient, discriminator="dos-bot"):
"""A bot that simulates a Denial of Service attack."""
class ConfigSchema(DatabaseClient.ConfigSchema):
"""ConfigSchema for DoSBot."""
type: str = "DoSBot"
type: str = "dos-bot"
target_ip_address: Optional[IPV4Address] = None
target_port: Port = PORT_LOOKUP["POSTGRES_SERVER"]
payload: Optional[str] = None
@@ -72,7 +72,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.name = "dos-bot"
self.target_ip_address = self.config.target_ip_address
self.target_port = self.config.target_port
self.payload = self.config.payload

8
src/primaite/simulator/system/applications/red_applications/ransomware_script.py
View File

@@ -14,7 +14,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class RansomwareScript(Application, identifier="RansomwareScript"):
class RansomwareScript(Application, discriminator="ransomware-script"):
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
:ivar payload: The attack stage query payload. (Default ENCRYPT)
@@ -23,7 +23,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for RansomwareScript."""
type: str = "RansomwareScript"
type: str = "ransomware-script"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "ENCRYPT"
@@ -38,7 +38,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
"Payload String for the payload stage"
def __init__(self, **kwargs):
kwargs["name"] = "RansomwareScript"
kwargs["name"] = "ransomware-script"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -63,7 +63,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the Ransomware Script."""
db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
db_client: DatabaseClient = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

8
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -23,7 +23,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebBrowser(Application, identifier="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
"""
Represents a web browser in the simulation environment.
@@ -33,7 +33,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for WebBrowser."""
type: str = "WebBrowser"
type: str = "web-browser"
target_url: Optional[str] = None
config: "WebBrowser.ConfigSchema" = Field(default_factory=lambda: WebBrowser.ConfigSchema())
@@ -48,7 +48,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
"""Keep a log of visited websites and information about the visit, such as response code."""
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["name"] = "web-browser"
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
@@ -108,7 +108,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
return False
# get the IP address of the domain name via DNS
dns_client: DNSClient = self.software_manager.software.get("DNSClient")
dns_client: DNSClient = self.software_manager.software.get("dns-client")
domain_exists = dns_client.check_domain_exists(target_domain=parsed_url.hostname)
# if domain does not exist, the request fails

6
src/primaite/simulator/system/core/software_manager.py
View File

@@ -60,12 +60,12 @@ class SoftwareManager:
@property
def arp(self) -> "ARP":
"""Provides access to the ARP service instance, if installed."""
return self.software.get("ARP") # noqa
return self.software.get("arp") # noqa
@property
def icmp(self) -> "ICMP":
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
return self.software.get("icmp") # noqa
def get_open_ports(self) -> List[Port]:
"""
@@ -244,7 +244,7 @@ class SoftwareManager:
:param session: The transport session the payload originates from.
"""
if payload.__class__.__name__ == "PortScanPayload":
self.software.get("NMAP").receive(payload=payload, session_id=session_id)
self.software.get("nmap").receive(payload=payload, session_id=session_id)
return
main_receiver = self.port_protocol_mapping.get((port, protocol), None)
if main_receiver:

6
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -15,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class ARP(Service, identifier="ARP"):
class ARP(Service, discriminator="arp"):
"""
The ARP (Address Resolution Protocol) Service.
@@ -26,14 +26,14 @@ class ARP(Service, identifier="ARP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ARP."""
type: str = "ARP"
type: str = "arp"
config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema())
arp: Dict[IPV4Address, ARPEntry] = {}
def __init__(self, **kwargs):
kwargs["name"] = "ARP"
kwargs["name"] = "arp"
kwargs["port"] = PORT_LOOKUP["ARP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

14
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DatabaseService(Service, identifier="DatabaseService"):
class DatabaseService(Service, discriminator="database-service"):
"""
A class for simulating a generic SQL Server service.
@@ -29,12 +29,12 @@ class DatabaseService(Service, identifier="DatabaseService"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DatabaseService."""
type: str = "DatabaseService"
type: str = "database-service"
backup_server_ip: Optional[IPv4Address] = None
db_password: Optional[str] = None
"""Password that needs to be provided by clients if they want to connect to the DatabaseService."""
config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DatabaseService.ConfigSchema())
backup_server_ip: IPv4Address = None
"""IP address of the backup server."""
@@ -46,7 +46,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
"""File name of latest backup."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseService"
kwargs["name"] = "database-service"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@@ -70,7 +70,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
"""
super().install()
if not self.parent.software_manager.software.get("FTPClient"):
if not self.parent.software_manager.software.get("ftp-client"):
self.parent.sys_log.info(f"{self.name}: Installing FTPClient to enable database backups")
self.parent.software_manager.install(FTPClient)
@@ -94,7 +94,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(
@@ -128,7 +128,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(

10
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -18,24 +18,22 @@ if TYPE_CHECKING:
_LOGGER = getLogger(__name__)
class DNSClient(Service, identifier="DNSClient"):
class DNSClient(Service, discriminator="dns-client"):
"""Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
dns_server: Optional[IPv4Address] = None
type: str = "dns-client"
dns_server: Optional[IPV4Address] = None
"The DNS Server the client sends requests to."
config: ConfigSchema = Field(default_factory=lambda: DNSClient.ConfigSchema())
dns_cache: Dict[str, IPv4Address] = {}
"A dict of known mappings between domain/URLs names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSClient"
kwargs["name"] = "dns-client"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

8
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -14,22 +14,22 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DNSServer(Service, identifier="DNSServer"):
class DNSServer(Service, discriminator="dns-server"):
"""Represents a DNS Server as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSServer."""
type: str = "DNSServer"
type: str = "dns-server"
domain_mapping: dict = {}
config: "DNSServer.ConfigSchema" = Field(default_factory=lambda: DNSServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DNSServer.ConfigSchema())
dns_table: Dict[str, IPv4Address] = {}
"A dict of mappings between domain names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSServer"
kwargs["name"] = "dns-server"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

10
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPClient(FTPServiceABC, identifier="FTPClient"):
class FTPClient(FTPServiceABC, discriminator="ftp-client"):
"""
A class for simulating an FTP client service.
@@ -26,15 +26,15 @@ class FTPClient(FTPServiceABC, identifier="FTPClient"):
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"
type: str = "ftp-client"
config: ConfigSchema = Field(default_factory=lambda: FTPClient.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["name"] = "ftp-client"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

9
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -12,7 +12,7 @@ from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPServer(FTPServiceABC, identifier="FTPServer"):
class FTPServer(FTPServiceABC, discriminator="ftp-server"):
"""
A class for simulating an FTP server service.
@@ -23,15 +23,14 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"):
class ConfigSchema(FTPServiceABC.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
type: str = "ftp-server"
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
config: ConfigSchema = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["name"] = "ftp-server"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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