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

#3111 & initial notebook fixes.

Browse Source
This commit is contained in:
Archer Bowen
2025-03-10 16:19:54 +00:00
parent b467292edd
commit b4df227af3
7 changed files with 287 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/primaite/notebooks/Action-masking.ipynb
View File

@@ -8,7 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"PrimAITE environments support action masking. The action mask shows which of the agent's actions are applicable with the current environment state. For example, a node can only be turned on if it is currently turned off."
"PrimAITE environments support action masking. The action mask shows which of the agent's actions are applicable with the current environment state. For example, a node can only be turned on if it is currently turned off. Please refer to the action masking configuration user guide page for more information."
]
},
{
@@ -151,7 +151,10 @@
"metadata": {},
"source": [
"## Action masking with MARL in Ray RLLib\n",
"Each agent has their own action mask, this is useful if the agents have different action spaces."
"\n",
"Each agent has their own action mask which is useful for multi-agent environments where each agent are configured with different action spaces.\n",
"\n",
"The code snippets below demonstrate how users can use multiple agents with action masks using the [UC2 MARL example config](./Training-an-RLLIB-MARL-System.ipynb)."
]
},
{

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

@@ -47,7 +47,7 @@
"\n",
"This notebook uses the same network setup as UC2. Please refer to the main [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb).\n",
"\n",
"However, this notebook replaces the red agent used in UC2 with a custom proxy red agent built for this notebook."
"However, this notebook replaces the red agent used in UC2 with a custom proxy red agent built for this notebook. This allows us to control the red agent manually via `env.step()` in the same way that other notebooks use the blue agent."
]
},
{
@@ -159,9 +159,9 @@
"## **Notebook Setup** | Network Prerequisites\n",
"\n",
"Before the Red Agent is able to perform any C2 specific actions, the C2 Server needs to be installed and run.\n",
"This is because in higher fidelity environments (and the real-world) a C2 server would not be accessible by a private network blue agent and the C2 Server would already be in place before the an adversary (Red Agent) starts.\n",
"This is because in higher fidelity environments (and the real-world) a C2 server would not be accessible by a private network blue agent and the C2 Server would already be in place before an adversary (red agent) began it's direct attack.\n",
"\n",
"The cells below install and run the C2 Server on client_1 directly via the simulation API."
"The cells below install and run the C2 Server on **client_1** directly via the simulation API."
]
},
{
@@ -192,26 +192,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Beacon Actions | node_application_install\n",
"### **Command and Control** | C2 Beacon Actions | `node-application-install`\n",
"\n",
"The custom proxy red agent defined at the start of this notebook has been configured to install the C2 Beacon as action ``1`` in it's action map. \n",
"\n",
"The below yaml snippet shows all the relevant agent options for this action:\n",
"The custom proxy red agent defined at the start of this notebook has been configured to install the C2 Beacon as action ``1`` in it's action map: \n",
"\n",
"```yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: c2-beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: node-application-install \n",
" action: node-application-install\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
"```"
]
@@ -231,11 +222,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Beacon Actions | configure_c2_beacon \n",
"### **Command and Control** | C2 Beacon Actions | `configure-c2-beacon` \n",
"\n",
"The custom proxy red agent defined at the start of this notebook can configure the C2 Beacon via action ``2`` in it's action map. \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"The custom proxy red agent defined at the start of this notebook can configure the C2 Beacon via action ``2`` in it's action map:\n",
"\n",
"```yaml\n",
"\n",
@@ -269,11 +258,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Beacon Actions | node_application_execute\n",
"### **Command and Control** | C2 Beacon Actions | ``node-application-execute``\n",
"\n",
"The final action is ``node-application-execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"The final action is ``node-application-execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map:\n",
"\n",
"```yaml\n",
" action_space:\n",
@@ -321,18 +308,16 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2-server-terminal-command\n",
"### **Command and Control** | C2 Server Actions | `c2-server-terminal-command`\n",
"\n",
"The C2 Server's terminal action: ``c2-server-terminal-command`` is indexed at ``4`` in it's action map. \n",
"\n",
"This action leverages the terminal service that is installed by default on all nodes to grant red agents a lot more configurability. If you're unfamiliar with terminals then it's recommended that you refer to the ``Terminal Processing`` notebook.\n",
"This action leverages the terminal service that is installed by default on all nodes to grant red agents a lot more configurability. If you're unfamiliar with terminals then it's recommended that you refer to the [Terminal-Processing notebook](./Terminal-Processing.ipynb) for more info.\n",
"\n",
"It's worth noting that an additional benefit a red agent has when using the terminal service via the C2 Server is that you can execute multiple commands in one action. \n",
"\n",
"In this notebook, the ``c2-server-terminal-command`` is used to install a RansomwareScript application on the ``web_server`` node.\n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"``` yaml\n",
" action_space:\n",
" action_map:\n",
@@ -376,11 +361,9 @@
"source": [
"### **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",
"Another action offered by the C2 server grants a red agent the ability to configure the `RansomwareScript` via the C2 Server rather than the node directly.\n",
"\n",
"This action is indexed as action ``5``.\n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"This action is indexed as action ``5``\n",
"\n",
"``` yaml\n",
" action_space:\n",
@@ -418,27 +401,25 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2-server-data-exfiltrate\n",
"### **Command and Control** | C2 Server Actions | ``c2-server-data-exfiltrate``\n",
"\n",
"The second to last action available is the ``c2-server-data-exfiltrate`` which is indexed as action ``6`` in the action map.\n",
"\n",
"This action can be used to exfiltrate a target file on a remote node to the C2 Beacon and the C2 Server's host file system via the ``FTP`` services.\n",
"\n",
"The below yaml snippet shows all the relevant agent options for this action:\n",
"\n",
"``` yaml\n",
" action_space:\n",
" action_map:\n",
" 6:\n",
" 6:\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" target_file_name: \"database.db\"\n",
" target_folder_name: \"database\"\n",
" exfiltration_folder_name: \"spoils\"\n",
" target_ip_address: \"192.168.1.14\"\n",
" username: \"admin\"\n",
" password: \"admin\"\n",
" target_ip_address: 192.168.1.14\n",
" username: admin\n",
" password: admin\n",
"\n",
"```"
]
@@ -476,21 +457,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2-server-ransomware-launch\n",
"### **Command and Control** | C2 Server Actions | ``c2-server-ransomware-launch``\n",
"\n",
"Finally, the last available action is for the c2-server-ransomware-launch to start the ransomware script installed on the same node as the C2 beacon.\n",
"\n",
"This action is indexed as action ``7``.\n",
"\n",
"\"The yaml snippet below shows all the relevant agent options for this action:\n",
"Finally, the last available action is for the c2-server-ransomware-launch to start the ransomware script installed previously:\n",
"\n",
"``` yaml\n",
" action_space:\n",
" action_map:\n",
" 7:\n",
" 7:\n",
" action: c2-server-ransomware-launch\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
"```\n"
]
},
@@ -770,11 +747,11 @@
"source": [
"#### **Command and Control** | OBS Impact | C2 Server | Terminal Command\n",
"\n",
"Using the C2 Server's ``TERMINAL`` command it is possible to install a ``RansomwareScript`` application onto the C2 Beacon's host.\n",
"As shown, we can use the C2 Server to install the ``RansomwareScript`` application onto the C2 Beacon's host.\n",
"\n",
"The below code cells perform this as well as capturing the OBS impacts.\n",
"\n",
"It's important to note that the ``TERMINAL`` command is not limited to just installing software."
"It's important to note that the ``c2-server-terminal-command`` action is not limited to just installing software as shown in the [Terminal-Processing](./Terminal-Processing.ipynb) notebook."
]
},
{
@@ -783,7 +760,7 @@
"metadata": {},
"outputs": [],
"source": [
"# Installing RansomwareScript via C2 Terminal Commands\n",
"# Installing RansomwareScript via C2 Terminal Commands (This is via the PrimAITE simulation rather than the c2-server-terminal-command action)\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"ransomware-script\"]],\n",
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
@@ -826,7 +803,7 @@
"source": [
"#### **Command and Control** | OBS Impact | C2 Server | Data Exfiltration\n",
"\n",
"Before encrypting the database.db file, the ``DATA_EXFILTRATION`` command can be used to copy the database.db file onto both the C2 Server and the C2 Beacon's file systems:"
"Before encrypting the database.db file, we can use the exfiltration command to copy the database.db file onto both the C2 Server and the C2 Beacon's file systems:"
]
},
{
@@ -988,7 +965,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The code cell below uses the custom blue agent defined at the start of this section perform a node_application_remove on the C2 beacon:"
"The code cell below uses the custom blue agent defined at the start of this section perform a `node_application_remove` on the C2 beacon:\n",
"\n",
"```yaml\n",
" action_space:\n",
" action_map:\n",
" 1:\n",
" action: node-application-remove\n",
" options:\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
"```"
]
},
{
@@ -1005,7 +992,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Which we can see after the effects of after stepping another timestep and looking at the web_servers software manager and the OBS differences."
"Which we can see after the effects of after stepping another timestep and looking at the **web_server**'s software manager and the OBS differences."
]
},
{
@@ -1083,7 +1070,16 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The code cell below uses the custom blue agent defined at the start of this section to perform a ``node_shut_down`` action on the web server."
"The code cell below uses the custom blue agent defined at the start of this section to perform a ``node_shut_down`` action on the web server.\n",
"\n",
"```yaml\n",
" action_space:\n",
" action_map:\n",
" 2:\n",
" action: node-shutdown\n",
" options:\n",
" node_name: web_server\n",
"```"
]
},
{
@@ -1173,7 +1169,27 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The code cell below uses the custom blue agent defined at the start of this section to perform a router_acl_add_rule on router 1."
"The code cell below uses the custom blue agent defined at the start of this section to perform a `router_acl_add_rule` on router 1.\n",
"\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_map:\n",
" 3:\n",
" action: router-acl-add-rule\n",
" options:\n",
" target_router: router_1\n",
" position: 1\n",
" permission: DENY\n",
" src_ip: 192.168.10.21\n",
" dst_ip: 192.168.1.12\n",
" src_port: HTTP\n",
" dst_port: HTTP\n",
" protocol_name: ALL\n",
" src_wildcard: 0.0.0.1\n",
" dst_wildcard: 0.0.0.1\n",
"```"
]
},
{
@@ -1236,7 +1252,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Because of the ACL rule the C2 beacon never received the ransomware installation and execute commands from the C2 server:"
"Because of the ACL rule the C2 beacon never receives the ransomware installation and execute commands from the C2 server:"
]
},
{
@@ -1254,7 +1270,7 @@
"metadata": {},
"outputs": [],
"source": [
"database_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"database-server\")\n",
"database_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"database_server\")\n",
"database_server.software_manager.file_system.show(full=True)"
]
},
@@ -1317,11 +1333,11 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | Configurability | C2 Server IP Address\n",
"### **Command and Control** | Configurability | `c2_server_ip_address`\n",
"\n",
"As with a majority of client and server based application configurations in primaite, the remote IP of a server must be supplied.\n",
"\n",
"In the case of the C2 Beacon, the C2 Server's IP address must be supplied before the C2 beacon will be able to perform any other actions (including ``APPLICATION EXECUTE``).\n",
"In the case of the C2 Beacon, the C2 Server's IP address must be supplied before the C2 beacon will be able to perform any other actions.\n",
"\n",
"If the network contains multiple C2 Servers then it's also possible to switch to a different C2 server mid-episode which is demonstrated in the below code cells."
]
@@ -1374,7 +1390,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Configuring the C2 Beacon to establish connection to the C2 Server on client_1 (192.168.10.21)"
"Configuring the C2 Beacon to establish connection to the C2 Server on **client_1** (*192.168.10.21)*"
]
},
{
@@ -1393,7 +1409,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Now reconfiguring the C2 Beacon to establish connection to the C2 Server on client_2 (192.168.10.22)"
"Now reconfiguring the C2 Beacon to establish connection to the C2 Server on **client_2** (*192.168.10.22*)"
]
},
{
@@ -1413,7 +1429,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"After six timesteps the client_1 server will recognise the C2 beacon's previous connection as dead and clear its connections. (This is dependent on the ``Keep Alive Frequency`` setting.)"
"After six timesteps the client_1 server will recognise the C2 beacon's previous connection as dead and clear its connections. (This is dependent on the ``keep_alive_frequency`` setting.)"
]
},
{
@@ -1432,9 +1448,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | Configurability | Keep Alive Frequency\n",
"### **Command and Control** | Configurability | `keep_alive_frequency`\n",
"\n",
"In order to confirm it's connection the C2 Beacon will send out a ``Keep Alive`` to the C2 Server and receive a keep alive back. \n",
"In order to confirm it's connection the C2 Beacon will send out a keep alive packet to the C2 Server and expects a keep alive sent back. \n",
"\n",
"By default, this occurs every 5 timesteps. However, this setting can be configured to be much more infrequent or as frequent as every timestep. \n",
"\n",
@@ -1480,7 +1496,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The next code cells capture the obs impact of the default Keep Alive Frequency which is 5 timesteps:"
"The next code cells capture the obs impact of the default `keep_alive_frequency` which is 5 timesteps:"
]
},
{
@@ -1536,7 +1552,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Demonstrating that the observation impacts of the Keep Alive can be seen on every timestep:"
"Now the keep alive's observation impacts can be seen on every timestep:"
]
},
{
@@ -1555,7 +1571,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Lastly, the keep_alive_frequency can also be used to configure the C2 Beacon to confirm connection less frequently. \n",
"Lastly, the `keep_alive_frequency` can also be used to configure the C2 Beacon to confirm connection less frequently. \n",
"\n",
"The code cells below demonstrate the impacts of changing the frequency rate to ``7`` timesteps."
]
@@ -1578,7 +1594,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | Configurability | Masquerade Port & Masquerade Protocol\n",
"### **Command and Control** | Configurability | `masquerade_port` & `masquerade_protocol`\n",
"\n",
"The final configurable options are ``Masquerade Port`` & ``Masquerade Protocol``. These options can be used to control the networking IP Protocol and Port the C2 traffic is currently using.\n",
"\n",
@@ -1586,7 +1602,7 @@
"\n",
"In primAITE, red agents can begin to simulate stealth behaviour by configuring C2 traffic to use different protocols mid episode or between episodes.\n",
"\n",
"Currently, red agent actions support the following port and protocol options:\n",
"Currently, the C2 Suite support the following port and protocol options:\n",
"\n",
"| Supported Ports | Supported Protocols |\n",
"|------------------|---------------------|\n",
@@ -1640,7 +1656,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"However, C2 Beacon can be configured to use UDP (``Masquerade Protocol``) and we can also configure the C2 Beacon to use a different Port (``Masquerade Port``) for example ``DNS``. "
"But we can use the aforementioned options to configure the C2 Beacon to use different protocol and port such as a `UDP` and `DNS`. "
]
},
{
@@ -1677,7 +1693,6 @@
"metadata": {},
"outputs": [],
"source": [
"\n",
"env.game.agents[\"CustomC2Agent\"].show_history()"
]
}

34
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb
View File

@@ -8,11 +8,11 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
"This notebook will go over some examples of how UC2 red agent behaviour can be varied by changing its configuration parameters.\n",
"\n",
"First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
"\n",
"*(For a full explanation of the Data Manipulation scenario, check out the data manipulation scenario notebook)*"
"*For a full explanation of the Data Manipulation scenario (also known as UC2), check out the [data manipulation scenario notebook](./Data-Manipulation-E2E-Demonstration.ipynb)*"
]
},
{
@@ -175,22 +175,20 @@
" network:\n",
" nodes:\n",
" - hostname: client_1\n",
" type: computer\n",
" ip_address: 192.168.10.21\n",
" subnet_mask: 255.255.255.0\n",
" default_gateway: 192.168.10.1\n",
" \n",
" # \n",
" applications:\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
" data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
" payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
" server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" - type: database-client # Database client must be installed in order for DataManipulationBot to function\n",
" options:\n",
" db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" type: computer\n",
" ip_address: 192.168.10.21\n",
" subnet_mask: 255.255.255.0\n",
" default_gateway: 192.168.10.1\n",
" applications:\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
" data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
" payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
" server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" - type: database-client # Database client must be installed in order for DataManipulationBot to function\n",
" options:\n",
" db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
"```"
]
},

10
src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -90,21 +90,21 @@
"source": [
"## Scripted agents:\n",
"### Red\n",
"The red agent sits on a client and uses an application called DataManipulationBot whose sole purpose is to send a DELETE query to the database.\n",
"The red agent sits on a client and uses an application called **DataManipulationBot** whose sole purpose is to send a DELETE query to the database.\n",
"The red agent can choose one of two action each timestep:\n",
"1. do nothing\n",
"2. execute the data manipulation application\n",
"The schedule for selecting when to execute the application is controlled by three parameters:\n",
"- start time\n",
"- frequency\n",
"- variance\n",
" - start time\n",
" - frequency\n",
" - variance\n",
"\n",
"Attacks start at a random timestep between (start_time - variance) and (start_time + variance). After each attack, another is attempted after a random delay between (frequency - variance) and (frequency + variance) timesteps.\n",
"\n",
"The data manipulation app itself has an element of randomness because the attack has a probability of success. The default is 0.8 to succeed with the port scan step and 0.8 to succeed with the attack itself.\n",
"Upon a successful attack, the database file becomes corrupted which incurs a negative reward for the RL defender.\n",
"\n",
"The red agent does not use information about the state of the network to decide its action.\n",
"*The red agent does not use information about the state of the network to decide its action. For a more complex red agent please refer to the [UC7 scenario](./UC7-E2E-Demo.ipynb)*\n",
"\n",
"### Green\n",
"The green agents use the web browser application to send requests to the web server. The schedule of each green agent is currently random, it will do nothing 30% of the time, send a web request 60% of the time, and send a db status check 10% of the time.\n",