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

#3075: Update notebooks to reflect extension changes.

Browse Source
This commit is contained in:
Nick Todd
2025-01-27 17:08:33 +00:00
parent 8eeba691a9
commit f4b73057d2
10 changed files with 116 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -22,13 +22,14 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 4,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
"from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.config.load import data_manipulation_config_path\n", "from primaite.config.load import data_manipulation_config_path\n",
"from prettytable import PrettyTable" "from prettytable import PrettyTable\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n"
] ]
}, },
{ {
@@ -103,7 +104,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 9,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -116,7 +117,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 10,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -156,7 +157,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 12,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -167,7 +168,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 13,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -204,7 +205,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -218,7 +219,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.11" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,

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

@@ -16,6 +16,15 @@
"execution_count": null, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [ "source": [
"# Imports\n", "# Imports\n",
"import yaml\n", "import yaml\n",
@@ -27,7 +36,8 @@
"from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import C2Command\n", "from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import C2Command\n",
"from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n", "from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n", "from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server" "from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent"
] ]
}, },
{ {
@@ -43,7 +53,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 6,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -1763,7 +1773,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },

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

@@ -26,7 +26,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 2,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -36,12 +36,13 @@
"from primaite.game.agent.interface import AgentHistoryItem\n", "from primaite.game.agent.interface import AgentHistoryItem\n",
"from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.session.environment import PrimaiteGymEnv\n",
"import yaml\n", "import yaml\n",
"from pprint import pprint" "from pprint import pprint\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent, data_manipulation_bot"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 3,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -68,7 +69,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 5,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -79,7 +80,8 @@
" if red_action == 'do_nothing':\n", " if red_action == 'do_nothing':\n",
" red_str = 'DO NOTHING'\n", " red_str = 'DO NOTHING'\n",
" elif red_action == 'node_application_execute':\n", " elif red_action == 'node_application_execute':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n", " client = \"client 1\" if red_info.parameters['node_name'] == 0 else \"client 2\"\n",
"\n",
" red_str = f\"ATTACK from {client}\"\n", " red_str = f\"ATTACK from {client}\"\n",
" return red_str" " return red_str"
] ]
@@ -243,7 +245,8 @@
"outputs": [], "outputs": [],
"source": [ "source": [
"change = yaml.safe_load(\"\"\"\n", "change = yaml.safe_load(\"\"\"\n",
"start_settings:\n", " possible_start_nodes: [client_1, client_2]\n",
" target_application: DataManipulationBot\n",
" start_step: 25\n", " start_step: 25\n",
" frequency: 20\n", " frequency: 20\n",
" variance: 0\n", " variance: 0\n",
@@ -253,7 +256,9 @@
" cfg = yaml.safe_load(f)\n", " cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n", " for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n", " if agent['ref'] == \"data_manipulation_attacker\":\n",
" print(f\"{agent['agent_settings']=}\")\n",
" agent['agent_settings'] = change\n", " agent['agent_settings'] = change\n",
" print(f\"{agent['agent_settings']=}\")\n",
"\n", "\n",
"env = PrimaiteGymEnv(env_config = cfg)\n", "env = PrimaiteGymEnv(env_config = cfg)\n",
"env.reset()\n", "env.reset()\n",
@@ -310,9 +315,17 @@
"outputs": [], "outputs": [],
"source": [ "source": [
"change = yaml.safe_load(\"\"\"\n", "change = yaml.safe_load(\"\"\"\n",
"# TODO:\n", " action_space:\n",
" action_map:\n",
" 0:\n",
" action: do_nothing\n",
" options: {}\n",
" 1:\n",
" action: node_application_execute\n",
" options:\n",
" node_name: client_1\n",
" application_name: DataManipulationBot\n",
"\"\"\")\n", "\"\"\")\n",
"#TODO 2869 fix\n",
"\n", "\n",
"with open(data_manipulation_config_path(), 'r') as f:\n", "with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n", " cfg = yaml.safe_load(f)\n",
@@ -438,7 +451,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -452,7 +465,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.11" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,

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

@@ -382,7 +382,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 2,
"metadata": { "metadata": {
"tags": [] "tags": []
}, },
@@ -394,7 +394,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 3,
"metadata": { "metadata": {
"tags": [] "tags": []
}, },
@@ -405,7 +405,8 @@
"from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.game.agent.interface import AgentHistoryItem\n", "from primaite.game.agent.interface import AgentHistoryItem\n",
"import yaml\n", "import yaml\n",
"from pprint import pprint\n" "from pprint import pprint\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent, data_manipulation_bot\n"
] ]
}, },
{ {
@@ -429,9 +430,9 @@
" cfg = yaml.safe_load(f)\n", " cfg = yaml.safe_load(f)\n",
" # set success probability to 1.0 to avoid rerunning cells.\n", " # set success probability to 1.0 to avoid rerunning cells.\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n", " cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n", " cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n", " cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n", " cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['port_scan_p_of_success'] = 1.0\n",
" # don't flatten observations so that we can see what is going on\n", " # don't flatten observations so that we can see what is going on\n",
" cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n", " cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n",
"\n", "\n",
@@ -450,7 +451,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 5,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -461,7 +462,7 @@
" if red_action == 'do_nothing':\n", " if red_action == 'do_nothing':\n",
" red_str = 'DO NOTHING'\n", " red_str = 'DO NOTHING'\n",
" elif red_action == 'node_application_execute':\n", " elif red_action == 'node_application_execute':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n", " client = \"client 1\" if red_info.parameters['node_name'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n", " red_str = f\"ATTACK from {client}\"\n",
" return red_str" " return red_str"
] ]

6
src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb
View File

@@ -32,6 +32,8 @@
"from primaite.session.environment import PrimaiteGymEnv\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.computer import Computer\n",
"from notebook.services.config import ConfigManager\n", "from notebook.services.config import ConfigManager\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"\n",
"\n", "\n",
"cm = ConfigManager().update('notebook', {'limit_output': 50}) # limit output lines to 50 - for neatness\n", "cm = ConfigManager().update('notebook', {'limit_output': 50}) # limit output lines to 50 - for neatness\n",
"\n", "\n",
@@ -157,7 +159,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": "venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -171,7 +173,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.8" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,

15
src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb
View File

@@ -54,6 +54,15 @@
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": { "metadata": {
"tags": [] "tags": []
}, },
@@ -68,7 +77,7 @@
"from primaite.simulator.network.hardware.nodes.host.server import Server\n", "from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.system.applications.database_client import DatabaseClient\n", "from primaite.simulator.system.applications.database_client import DatabaseClient\n",
"from primaite.simulator.system.applications.web_browser import WebBrowser\n", "from primaite.simulator.system.applications.web_browser import WebBrowser\n",
"from primaite.simulator.system.services.database.database_service import DatabaseService" "from primaite.simulator.system.services.database.database_service import DatabaseService\n"
] ]
}, },
{ {
@@ -80,7 +89,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 3,
"metadata": { "metadata": {
"tags": [] "tags": []
}, },
@@ -103,7 +112,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 4,
"metadata": { "metadata": {
"tags": [] "tags": []
}, },

20
src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb
View File

@@ -40,6 +40,7 @@
"import ray\n", "import ray\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n", "from ray.rllib.algorithms.ppo import PPOConfig\n",
"from primaite.session.ray_envs import PrimaiteRayMARLEnv\n", "from primaite.session.ray_envs import PrimaiteRayMARLEnv\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"\n", "\n",
"with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n", "with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n",
" cfg = yaml.safe_load(f)\n", " cfg = yaml.safe_load(f)\n",
@@ -56,7 +57,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 2,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -102,7 +103,20 @@
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"*** SIGTERM received at time=1737996337 on cpu 7 ***\n",
"PC: @ 0x7f3649b0fe2e (unknown) epoll_wait\n",
" @ 0x7f3649a2c520 (unknown) (unknown)\n",
"[2025-01-27 16:45:37,381 E 117142 117142] logging.cc:440: *** SIGTERM received at time=1737996337 on cpu 7 ***\n",
"[2025-01-27 16:45:37,381 E 117142 117142] logging.cc:440: PC: @ 0x7f3649b0fe2e (unknown) epoll_wait\n",
"[2025-01-27 16:45:37,381 E 117142 117142] logging.cc:440: @ 0x7f3649a2c520 (unknown) (unknown)\n"
]
}
],
"source": [ "source": [
"eval = algo.evaluate()" "eval = algo.evaluate()"
] ]
@@ -110,7 +124,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": "venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },

23
src/primaite/notebooks/Training-an-RLLib-Agent.ipynb
View File

@@ -32,6 +32,8 @@
"from primaite.session.ray_envs import PrimaiteRayEnv\n", "from primaite.session.ray_envs import PrimaiteRayEnv\n",
"import ray\n", "import ray\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n", "from ray.rllib.algorithms.ppo import PPOConfig\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"\n",
"\n", "\n",
"# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n", "# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n",
"# to copy the files to your user data path.\n", "# to copy the files to your user data path.\n",
@@ -50,7 +52,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 3,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -96,7 +98,20 @@
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"*** SIGTERM received at time=1737996055 on cpu 1 ***\n",
"PC: @ 0x7f6e254a6e2e (unknown) epoll_wait\n",
" @ 0x7f6e253c3520 (unknown) (unknown)\n",
"[2025-01-27 16:40:55,343 E 114171 114171] logging.cc:440: *** SIGTERM received at time=1737996055 on cpu 1 ***\n",
"[2025-01-27 16:40:55,343 E 114171 114171] logging.cc:440: PC: @ 0x7f6e254a6e2e (unknown) epoll_wait\n",
"[2025-01-27 16:40:55,344 E 114171 114171] logging.cc:440: @ 0x7f6e253c3520 (unknown) (unknown)\n"
]
}
],
"source": [ "source": [
"eval = algo.evaluate()" "eval = algo.evaluate()"
] ]
@@ -104,7 +119,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -118,7 +133,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.11" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,

15
src/primaite/notebooks/Using-Episode-Schedules.ipynb
View File

@@ -48,6 +48,7 @@
"from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite import PRIMAITE_PATHS\n", "from primaite import PRIMAITE_PATHS\n",
"from prettytable import PrettyTable\n", "from prettytable import PrettyTable\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent, data_manipulation_bot\n",
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/scenario_with_placeholders\"" "scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/scenario_with_placeholders\""
] ]
}, },
@@ -409,21 +410,9 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": "venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
},
"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.11"
} }
}, },
"nbformat": 4, "nbformat": 4,

18
src/primaite/notebooks/multi-processing.ipynb
View File

@@ -29,7 +29,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 30,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -37,13 +37,13 @@
"from stable_baselines3 import PPO\n", "from stable_baselines3 import PPO\n",
"from stable_baselines3.common.utils import set_random_seed\n", "from stable_baselines3.common.utils import set_random_seed\n",
"from stable_baselines3.common.vec_env import SubprocVecEnv\n", "from stable_baselines3.common.vec_env import SubprocVecEnv\n",
"\n", "from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.session.environment import PrimaiteGymEnv\n" "from primaite.game.agent.scripted_agents import probabilistic_agent\n"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 31,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -52,7 +52,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 32,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -69,7 +69,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 33,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -90,7 +90,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 36,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -138,7 +138,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -152,7 +152,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.11" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,