Files
PrimAITE/src/primaite/notebooks/UC7-attack-variants.ipynb

587 lines
19 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# UC7 with Attack Variability\n",
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK"
]
},
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": [
"This notebook demonstrates the PrimAITE environment with the UC7 network laydown and multiple attack personas. The first attack persona is TAP001 which performs a ransomware attack against the database. The other one is TAP003 which is able to maliciously add ACL rules that block green pattern of life.\n",
"\n",
"The environment switches between these two attacks on a pre-defined schedule which is defined in the schedule.yaml file of the scenario folder."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup and Imports"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import yaml\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite import PRIMAITE_PATHS\n",
"from prettytable import PrettyTable\n",
"from deepdiff.diff import DeepDiff\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.network.hardware.nodes.network.router import Router\n",
"from primaite.simulator.system.services.dns.dns_server import DNSServer\n",
"from primaite.simulator.system.software import SoftwareHealthState\n",
"from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus\n",
"from primaite.simulator.network.hardware.nodes.network.switch import Switch\n",
"from primaite.simulator.system.applications.web_browser import WebBrowser\n",
"from primaite.simulator.network.container import Network\n",
"from primaite.simulator.system.services.service import ServiceOperatingState\n",
"from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState\n",
"from primaite.simulator.system.services.database.database_service import DatabaseService\n",
"from primaite.simulator.system.applications.database_client import DatabaseClient\n",
"from primaite.simulator.network.hardware.nodes.network.firewall import Firewall\n",
"from primaite.game.game import PrimaiteGame\n",
"from primaite.simulator.sim_container import Simulation\n",
"from primaite.config.load import load, _EXAMPLE_CFG\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.network.hardware.nodes.network.router import Router\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"\n",
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/uc7_multiple_attack_variants\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env = PrimaiteGymEnv(env_config=scenario_path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Schedule\n",
"\n",
"Let's print the schedule so that we can see which attack we can expect on each episode.\n",
"\n",
"On episodes 0-4, the TAP001 agent will be used, and on episodes 5-9, the TAP003 agent will be used. Then, the environment will alternate between the two. Furthermore, the TAP001 agent will alternate between starting at `ST_PROJ-A-PRV-PC-1`, `ST_PROJ-B-PRV-PC-2`, `ST_PROJ-C-PRV-PC-3`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(scenario_path / \"schedule.yaml\",'r') as f:\n",
" print(f.read())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP001 attack\n",
"\n",
"Let's first demonstrate the TAP001 attack. We will let the environment run for 30 steps and print out the red agent's actions.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#utils\n",
"def run_green_and_red_pol(num_steps):\n",
" for i in range(num_steps): # perform steps\n",
" env.step(0)\n",
"\n",
"def print_agent_actions_except_do_nothing(agent_name):\n",
" \"\"\"Get the agent's action history, filter out `do-nothing` actions, print relevant data in a table.\"\"\"\n",
" table = PrettyTable()\n",
" table.field_names = [\"Step\", \"Action\", \"Node\", \"Application\", \"Target IP\", \"Response\"]\n",
" print(f\"Episode: {env.episode_counter}, Actions for '{agent_name}':\")\n",
" for item in env.game.agents[agent_name].history:\n",
" if item.action == \"do-nothing\":\n",
" continue\n",
"\n",
" node, application, target_ip = \"N/A\", \"N/A\", \"N/A\",\n",
"\n",
" if item.action.startswith(\"node-nmap\"):\n",
" node = item.parameters['source_node']\n",
" application = \"nmap\"\n",
" target_ip = str(item.parameters['target_ip_address'])\n",
" target_ip = (target_ip[:25]+'...') if len(target_ip)>25 else target_ip # truncate long string\n",
"\n",
" elif item.action == \"router-acl-add-rule\":\n",
" node = item.parameters.get(\"router_name\")\n",
" elif item.action == \"node-send-remote-command\":\n",
" node = item.parameters.get(\"node_name\")\n",
" target_ip = item.parameters.get(\"remote_ip\")\n",
" application = item.parameters.get(\"command\")\n",
" elif item.action == \"node-session-remote-login\":\n",
" node = item.parameters.get(\"node_name\")\n",
" target_ip = item.parameters.get(\"remote_ip\")\n",
" application = \"user-manager\"\n",
" elif item.action.startswith(\"c2-server\"):\n",
" application = \"c2-server\"\n",
" node = item.parameters.get('node_name')\n",
" elif item.action == \"configure-c2-beacon\":\n",
" application = \"c2-beacon\"\n",
" node = item.parameters.get('node_name')\n",
"\n",
" else:\n",
" if (node_id := item.parameters.get('node_id')) is not None:\n",
" node = env.game.agents[agent_name].action_manager.node_names[node_id]\n",
" if (application_id := item.parameters.get('application_id')) is not None:\n",
" application = env.game.agents[agent_name].action_manager.application_names[node_id][application_id]\n",
" if (application_name := item.parameters.get('application_name')) is not None:\n",
" application = application_name\n",
"\n",
" table.add_row([item.timestep, item.action, node, application, target_ip, item.response.status])\n",
"\n",
" print(table)\n",
" print(\"(Any DONOTHING actions are omitted)\")\n",
"\n",
"def finish_episode_and_print_reward():\n",
" while env.game.step_counter < 128:\n",
" env.step(0)\n",
" print(f\"Total reward this episode: {env.agent.reward_function.total_reward:2f}\")\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"run_green_and_red_pol(110)\n",
"print_agent_actions_except_do_nothing(\"attacker\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_data_prv_srv_db: Server = env.game.simulation.network.get_node_by_hostname(\"ST_DATA-PRV-SRV-DB\")\n",
"st_data_prv_srv_db.file_system.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP001 Prevention"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The blue agent should be able to prevent the ransomware attack by blocking the red agent's access to the database. Let's run the environment until the observation space shows symptoms of the attack starting.\n",
"\n",
"Because we are in episode index 1, the red agent will use `ST-PROJ-A-PRV-PC-1` to start the attack. On step 25, the red agent installs `RansomwareScript`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"obs, reward, term, trunc, info = env.step(0)\n",
"for i in range(25): # we know that the ransomware install happens at step 25\n",
" old = obs\n",
" obs, reward, term, trunc, info = env.step(0)\n",
" new = obs\n",
"\n",
"diff = DeepDiff(old,new)\n",
"print(f\"Step {env.game.step_counter}\") # it's step 26 now because the step counter is incremented after the step\n",
"for d,v in diff.get('values_changed', {}).items():\n",
" print(f\"{d}: {v['old_value']} -> {v['new_value']}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that on HOST0, application index 1 has gone from `operating_status` 0 to 3, meaning there wasn't an application before, but now there is an application in the `INSTALLING` state. The blue agent should be able to detect this and block the red agent's access to the database. Action 43 will block `ST-PROJ-A-PRV-PC-1` from sending POSTGRES traffic to the DB server.\n",
"\n",
"If this were a different episode, it could have been `ST-PROJ-B-PRV-PC-2` or `ST-PROJ-C-PRV-PC-3` that are affected, and a different defensive action would be required."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(43)\n",
"env.step(45)\n",
"env.step(47)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_intra_prv_rt_cr: Router = env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-CR\")\n",
"st_intra_prv_rt_cr.acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_intra_prv_rt_cr.acl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now TAP001 is unable to locate the database!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print_agent_actions_except_do_nothing(\"attacker\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP003 attack\n",
"\n",
"Let's skip until episode 5 and demonstrate the TAP003 attack. We will let the environment run and print out the red agent's actions.\n",
"\n",
"By default, TAP003 will add the following rules:\n",
"\n",
"|Target Router | Impact |\n",
"|----------------------|--------|\n",
"|`ST_INTRA-PRV-RT-DR-1`| Blocks all `POSTGRES_SERVER` that arrives at the `ST_INTRA-PRV-RT-DR-1` router. This rule will prevent all ST_PROJ_* hosts from accessing the database (`ST_DATA-PRV-SRV-DB`).|\n",
"|`ST_INTRA-PRV-RT-CR`| Blocks all `HTTP` traffic that arrives at the`ST_INTRA-PRV-RT-CR` router. This rule will prevent all SOME_TECH hosts from accessing the webserver (`ST-DMZ-PUB-SRV-WEB`)|\n",
"|`REM-PUB-RT-DR`| Blocks all `DNS` traffic that arrives at the `REM-PUB-RT-DR` router. This rule prevents any remote site works from accessing the DNS Server (`ISP-PUB-SRV-DNS`).|"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"while env.episode_counter < 5:\n",
" env.reset()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"run_green_and_red_pol(128)\n",
"print_agent_actions_except_do_nothing(\"attacker\")\n",
"obs, reward, term, trunc, info = env.step(0); # one more step so we can capture the value of `obs`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The agent selected to add ACL rules that will prevent green pattern of life by blocking a variety of different traffic. This has a negative impact on reward. Let's view the ACL list on the affected router."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-DR-1\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-CR\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that at indices 1-5, there are ACL rules that block all traffic. The blue agent can see this rule in the `ROUTERS` part of the observation space.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER0']['ACL'][1]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER1']['ACL'][1]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER2']['ACL'][1]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Preventing TAP003 attack\n",
"\n",
"The blue agent can prevent the red agent from adding ACL rules. TAP003 relies on connecting to the router via SSH, and sending remote ACL_ADDRULE requests. The blue agent can prevent this by pre-emptively changing the admin password on the affected routers or by blocking SSH traffic between the red agent's starting node and the target routers."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"obs, reward, term, trunc, info = env.step(0)\n",
"old = obs\n",
"for i in range(128): \n",
" obs, reward, term, trunc, info = env.step(0)\n",
" new = obs\n",
"\n",
"diff = DeepDiff(old,new)\n",
"print(f\"Step {env.game.step_counter}\") # it's the next step now because the step counter is incremented after the step\n",
"for d,v in diff.get('values_changed', {}).items():\n",
" print(f\"{d}: {v['old_value']} -> {v['new_value']}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By printing the reward of each individual agent, we will see what green agents are affected the most. Of course, these green rewards count towards the blue reward so ultimately the blue agent should learn to remove the ACL rule."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The most effective option that the blue agent has against TAP003 is to prevent the red agent from ever adding the ACLs in the first place through blocking the SSH connection."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"env.step(51) # SSH Blocking ACL on ST-INRA-PRV-RT-R1\n",
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, another option the blue agent can take is to change the passwords of the different target routers that TAP003 will attack through the `NODE_ACCOUNTS_CHANGE_PASSWORD` action."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"env.step(50) # NODE_ACCOUNTS_CHANGE_PASSWORD | ST_INTRA-prv-rt-cr\n",
"env.step(52) # NODE_ACCOUNTS_CHANGE_PASSWORD | ST_INTRA-prv-rt-dr-1\n",
"env.step(54) # NODE_ACCOUNTS_CHANGE_PASSWORD | rem-pub-rt-dr\n",
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lastly, the blue agent can remedy the impacts of TAP003 through removing the malicious ACLs that TAP003 adds."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"\n",
"# Allow TAP003 to add it's malicious rules\n",
"for _ in range(45):\n",
" env.step(0)\n",
"\n",
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-CR\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-DR-1\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(44) # ROUTER_ACL_REMOVERULE | ST_INTRA-prv-rt-cr\n",
"env.step(53) # ROUTER_ACL_REMOVERULE | ST_INTRA-prv-rt-dr-1\n",
"env.step(55) # ROUTER_ACL_REMOVERULE | rem-pub-rt-dr"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-CR\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"ST_INTRA-PRV-RT-DR-1\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "milpac_venv",
"language": "python",
"name": "milpac_venv"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}