From e45c2742ccd5037fc0fe20225dc1faaecbda4ad3 Mon Sep 17 00:00:00 2001 From: Archer Bowen Date: Wed, 26 Feb 2025 16:59:35 +0000 Subject: [PATCH] #2925 Responding to a few more pr comments (general notebook clean up and improvement). --- .../notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb | 46 +++++++++++++++---- src/primaite/notebooks/UC7-Training.ipynb | 9 ---- .../notebooks/UC7-network_connectivity.ipynb | 25 ++-------- 3 files changed, 41 insertions(+), 39 deletions(-) diff --git a/src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb b/src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb index 8ec96da4..983890e7 100644 --- a/src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb +++ b/src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb @@ -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()" diff --git a/src/primaite/notebooks/UC7-Training.ipynb b/src/primaite/notebooks/UC7-Training.ipynb index f15084ac..ddaa6844 100644 --- a/src/primaite/notebooks/UC7-Training.ipynb +++ b/src/primaite/notebooks/UC7-Training.ipynb @@ -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, diff --git a/src/primaite/notebooks/UC7-network_connectivity.ipynb b/src/primaite/notebooks/UC7-network_connectivity.ipynb index 1a2415ac..8aef5465 100644 --- a/src/primaite/notebooks/UC7-network_connectivity.ipynb +++ b/src/primaite/notebooks/UC7-network_connectivity.ipynb @@ -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,