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

Merge branch 'dev' into bugfix/2299-check_hash_function_corrupts_files_and_folders

Browse Source
This commit is contained in:
Nick Todd
2024-04-23 16:40:02 +01:00
parent 90224960e5 a60e8b2203
commit 3eea70f7d8
140 changed files with 9846 additions and 5182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -22,6 +22,7 @@
"# Imports\n",
"\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"import yaml\n",
"from pprint import pprint"
@@ -62,12 +63,12 @@
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info[0]\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'NODE_APPLICATION_EXECUTE':\n",
" client = \"client 1\" if red_info[1]['node_id'] == 0 else \"client 2\"\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
@@ -361,7 +362,7 @@
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" if node['hostname'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
@@ -406,7 +407,7 @@
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" if node['hostname'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",

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

@@ -208,7 +208,7 @@
"|--|--|\n",
"|0|UNUSED|\n",
"|1|GOOD|\n",
"|2|PATCHING|\n",
"|2|FIXING|\n",
"|3|COMPROMISED|\n",
"|4|OVERWHELMED|\n",
"\n",
@@ -352,7 +352,7 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {
"tags": []
},
@@ -364,7 +364,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {
"tags": []
},
@@ -373,7 +373,7 @@
"# Imports\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.game.game import PrimaiteGame\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"import yaml\n",
"from pprint import pprint\n"
]
@@ -389,162 +389,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2024-03-13 16:52:48,201: Resetting environment, episode 0, avg. reward: 0.0\n",
"2024-03-13 16:52:48,205: Saving agent action log to C:\\Users\\NickTodd\\primaite\\3.0.0b6\\sessions\\2024-03-13\\16-52-48\\agent_actions\\episode_0.json\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"env created successfully\n",
"{'ACL': {1: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 0,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 2: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 1,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 3: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 2,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 4: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 3,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 5: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 4,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 6: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 5,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 7: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 6,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 8: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 7,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 9: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 8,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 10: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 9,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0}},\n",
" 'ICS': 0,\n",
" 'LINKS': {1: {'PROTOCOLS': {'ALL': 1}},\n",
" 2: {'PROTOCOLS': {'ALL': 1}},\n",
" 3: {'PROTOCOLS': {'ALL': 1}},\n",
" 4: {'PROTOCOLS': {'ALL': 1}},\n",
" 5: {'PROTOCOLS': {'ALL': 1}},\n",
" 6: {'PROTOCOLS': {'ALL': 1}},\n",
" 7: {'PROTOCOLS': {'ALL': 1}},\n",
" 8: {'PROTOCOLS': {'ALL': 1}},\n",
" 9: {'PROTOCOLS': {'ALL': 1}},\n",
" 10: {'PROTOCOLS': {'ALL': 0}}},\n",
" 'NODES': {1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 2: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 3: {'FOLDERS': {1: {'FILES': {1: {'health_status': 1}},\n",
" 'health_status': 1}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 4: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 5: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 6: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 7: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 1},\n",
" 2: {'NMNE': {'inbound': 0, 'outbound': 0},\n",
" 'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}}\n"
]
}
],
"outputs": [],
"source": [
"# create the env\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
@@ -565,20 +412,9 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"res = FileSystemItemHealthStatus.GOOD\n",
"res = FileSystemItemHealthStatus.GOOD\n",
"res = FileSystemItemHealthStatus.COMPROMISED\n",
"res = FileSystemItemHealthStatus.COMPROMISED\n"
]
}
],
"outputs": [],
"source": [
"# Test NODE_FOLDER_CHECKHASH\n",
"res = env.game.simulation.network.get_node_by_hostname('database_server').file_system.get_folder(folder_name = 'database').health_status\n",
@@ -618,12 +454,12 @@
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info[0]\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'NODE_APPLICATION_EXECUTE':\n",
" client = \"client 1\" if red_info[1]['node_id'] == 0 else \"client 2\"\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
@@ -643,7 +479,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Now the reward is -1, let's have a look at blue agent's observation."
"Now the reward is -0.8, let's have a look at blue agent's observation."
]
},
{
@@ -704,9 +540,9 @@
"source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'].action}\" )\n",
"print(f\"Blue reward:{reward}\" )"
]
},
@@ -714,7 +550,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The patching takes two steps, so the reward hasn't changed yet. Let's do nothing for another timestep, the reward should improve.\n",
"The fixing takes two steps, so the reward hasn't changed yet. Let's do nothing for another timestep, the reward should improve.\n",
"\n",
"The reward will increase slightly as soon as the file finishes restoring. Then, the reward will increase to 1 when both green agents make successful requests.\n",
"\n",
@@ -727,9 +563,9 @@
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
"obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'][0]}\" )\n",
"print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user']}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user']}\" )\n",
"print(f\"Blue reward:{reward:.2f}\" )"
@@ -751,24 +587,26 @@
"outputs": [],
"source": [
"env.step(13) # Patch the database\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'][0]}, Blue reward:{reward:.2f}\" )\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"env.step(50) # Block client 1\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'][0]}, Blue reward:{reward:.2f}\" )\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"env.step(51) # Block client 2\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'][0]}, Blue reward:{reward:.2f}\" )\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
"\n",
"for step in range(30):\n",
"while abs(reward - 0.8) > 1e-5:\n",
" obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'][0]}, Blue reward:{reward:.2f}\" )"
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
" if env.game.step_counter > 10000:\n",
" break # make sure there's no infinite loop if something went wrong"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, even though the red agent executes an attack, the reward stays at 0.8."
"Now, even though the red agent executes an attack, the reward will stay at 0.8."
]
},
{
@@ -784,7 +622,7 @@
"metadata": {},
"outputs": [],
"source": [
"obs['ACL']"
"obs['NODES']['ROUTER0']"
]
},
{
@@ -800,13 +638,30 @@
"metadata": {},
"outputs": [],
"source": [
"if obs['NODES'][6]['NETWORK_INTERFACES'][1]['nmne']['outbound'] == 1:\n",
" # client 1 has NMNEs, let's unblock client 2\n",
" env.step(58) # remove ACL rule 6\n",
"elif obs['NODES'][7]['NETWORK_INTERFACES'][1]['nmne']['outbound'] == 1:\n",
" env.step(57) # remove ACL rule 5\n",
"else:\n",
" print(\"something went wrong, neither client has NMNEs\")"
"env.step(58) # Remove the ACL rule that blocks client 1\n",
"env.step(57) # Remove the ACL rule that blocks client 2\n",
"\n",
"tries = 0\n",
"while True:\n",
" tries += 1\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
"\n",
" if obs['NODES']['HOST5']['NICS'][1]['NMNE']['outbound'] == 1:\n",
" # client 1 has NMNEs, let's block it\n",
" obs, reward, terminated, truncated, info = env.step(50) # block client 1\n",
" print(\"blocking client 1\")\n",
" break\n",
" elif obs['NODES']['HOST6']['NICS'][1]['NMNE']['outbound'] == 1:\n",
" # client 2 has NMNEs, so let's block it\n",
" obs, reward, terminated, truncated, info = env.step(51) # block client 2\n",
" print(\"blocking client 2\")\n",
" break\n",
" if tries>100:\n",
" print(\"Error: NMNE never increased\")\n",
" break\n",
"\n",
"env.step(13) # Patch the database\n",
"print()\n"
]
},
{
@@ -824,14 +679,14 @@
"source": [
"for step in range(30):\n",
" obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'][0]}, Blue reward:{reward:.2f}\" )"
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Reset the environment, you can rerun the other cells to verify that the attack works the same every episode."
"Reset the environment, you can rerun the other cells to verify that the attack works the same every episode. (except the red agent will move between `client_1` and `client_2`.)"
]
},
{

98
src/primaite/notebooks/Training-an-SB3-Agent.ipynb
View File

@@ -1,5 +1,21 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Training an SB3 Agent\n",
"\n",
"This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### First, we import the inital packages and read in our configuration file."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -27,7 +43,14 @@
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n"
" cfg = yaml.safe_load(f)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Using the given configuration, we generate the environment our agent will train in."
]
},
{
@@ -40,12 +63,10 @@
]
},
{
"cell_type": "code",
"execution_count": null,
"cell_type": "markdown",
"metadata": {},
"outputs": [],
"source": [
"from stable_baselines3 import PPO"
"Lets define training parameters for the agent."
]
},
{
@@ -54,7 +75,13 @@
"metadata": {},
"outputs": [],
"source": [
"model = PPO('MlpPolicy', gym)\n"
"from stable_baselines3 import PPO\n",
"\n",
"EPISODE_LEN = 128\n",
"NUM_EPISODES = 10\n",
"NO_STEPS = EPISODE_LEN * NUM_EPISODES\n",
"BATCH_SIZE = 32\n",
"LEARNING_RATE = 3e-4"
]
},
{
@@ -63,7 +90,14 @@
"metadata": {},
"outputs": [],
"source": [
"model.learn(total_timesteps=10)\n"
"model = PPO('MlpPolicy', gym, learning_rate=LEARNING_RATE, n_steps=NO_STEPS, batch_size=BATCH_SIZE, verbose=0, tensorboard_log=\"./PPO_UC2/\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"With the agent configured, let's train for our defined number of episodes."
]
},
{
@@ -72,7 +106,14 @@
"metadata": {},
"outputs": [],
"source": [
"model.save(\"deleteme\")"
"model.learn(total_timesteps=NO_STEPS)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's save the agent to a zip file that can be used in future evaluation."
]
},
{
@@ -80,7 +121,44 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
"source": [
"model.save(\"PrimAITE-PPO-example-agent\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we load the saved agent and run it in evaluation mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"eval_model = PPO(\"MlpPolicy\", gym)\n",
"eval_model = PPO.load(\"PrimAITE-PPO-example-agent\", gym)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, evaluate the agent."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from stable_baselines3.common.evaluation import evaluate_policy\n",
"\n",
"evaluate_policy(eval_model, gym, n_eval_episodes=10)"
]
}
],
"metadata": {
@@ -99,7 +177,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,