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

#2925 Responding to a few more pr comments (general notebook clean up and improvement).

Browse Source
This commit is contained in:
Archer Bowen
2025-02-26 16:59:35 +00:00
committed by Marek Wolan
parent eec73deae7
commit e45c2742cc
3 changed files with 41 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb
View File

@@ -66,26 +66,51 @@
"# Utility functions.\n",
"\n",
"def print_agent_actions_except_do_nothing(agent_name):\n",
" \"\"\"\n",
" Print an agent action provided it's not the DONOTHING action.\n",
"\n",
" :param agent_name: Name of agent (str).\n",
" \"\"\"\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\", \"Response\"]\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 != \"DONOTHING\":\n",
" node, application = 'unknown', 'unknown'\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",
" table.add_row([item.timestep, item.action, node, application, item.response.status])\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"
" print(\"(Any do-nothing actions are omitted)\")\n"
]
},
{
@@ -202,6 +227,7 @@
"with open(_EXAMPLE_CFG/\"uc7_config_tap003.yaml\", mode=\"r\") as uc7_config:\n",
" cfg = yaml.safe_load(uc7_config)\n",
" cfg[\"agents\"][33][\"agent_settings\"][\"flatten_obs\"] = False\n",
" cfg['io_settings']['save_agent_logs'] = True\n",
" cfg['io_settings']['save_sys_logs'] = True # Saving syslogs\n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"env.game.simulation.network.show()"

9
src/primaite/notebooks/UC7-Training.ipynb
View File

@@ -49,15 +49,6 @@
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/uc7_config.yaml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,

25
src/primaite/notebooks/UC7-network_connectivity.ipynb
View File

@@ -58,8 +58,11 @@
"metadata": {},
"outputs": [],
"source": [
"use_case_7_config = load(_EXAMPLE_CFG/\"uc7_config.yaml\")\n",
"env = PrimaiteGymEnv(env_config=use_case_7_config)"
"with open(file=_EXAMPLE_CFG/\"uc7_config.yaml\", mode=\"r\") as uc7_config:\n",
" cfg = yaml.safe_load(uc7_config)\n",
" cfg['io_settings']['save_sys_logs'] = True # Saving syslogs\n",
" cfg['io_settings']['save_agent_logs'] = True # Save agent logs\n",
"env = PrimaiteGymEnv(env_config=cfg)"
]
},
{
@@ -998,24 +1001,6 @@
"env.step(0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# from primaite.simulator.system.applications.application import ApplicationOperatingState\n",
"# from primaite.simulator.system.applications.database_client import database-client\n",
"# from primaite.simulator.system.services.service import ServiceOperatingState\n",
"\n",
"# home_pub_pc_1_dns_client: dns-client = home_pub_pc_1.software_manager.software[\"dns-client\"]\n",
"# home_pub_pc_1_database_client: database-client = home_pub_pc_1.software_manager.software[\"database-client\"]\n",
"\n",
"\n",
"# home_pub_pc_1_dns_client.operating_state = ServiceOperatingState.PAUSED\n",
"# home_pub_pc_1_database_client.operating_state = ApplicationOperatingState.CLOSED"
]
},
{
"cell_type": "code",
"execution_count": null,