From f54f278fca6e71d48c5bf375a03933b8db200891 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 21 Sep 2023 10:13:01 +0100 Subject: [PATCH 01/35] Initialise observations in agent interface --- sandbox.ipynb | 134 ++++++++++++++++++ src/primaite/game/__init__.py | 0 src/primaite/game/actor/__init__.py | 0 src/primaite/game/actor/actions.py | 21 +++ src/primaite/game/actor/interface.py | 32 +++++ src/primaite/game/actor/observations.py | 107 ++++++++++++++ src/primaite/game/actor/rewards.py | 20 +++ src/primaite/game/session.py | 6 + .../simulator/file_system/file_system.py | 16 ++- src/primaite/simulator/network/container.py | 4 +- 10 files changed, 333 insertions(+), 7 deletions(-) create mode 100644 sandbox.ipynb create mode 100644 src/primaite/game/__init__.py create mode 100644 src/primaite/game/actor/__init__.py create mode 100644 src/primaite/game/actor/actions.py create mode 100644 src/primaite/game/actor/interface.py create mode 100644 src/primaite/game/actor/observations.py create mode 100644 src/primaite/game/actor/rewards.py create mode 100644 src/primaite/game/session.py diff --git a/sandbox.ipynb b/sandbox.ipynb new file mode 100644 index 00000000..e7db5f4c --- /dev/null +++ b/sandbox.ipynb @@ -0,0 +1,134 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.network.networks import arcd_uc2_network\n", + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "net = arcd_uc2_network()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "random_node = list(net.nodes.keys())[0]\n", + "f = net.nodes[random_node].file_system.create_file(file_name=\"testfile\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "f.describe_state()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def test_file_observation():\n", + " from primaite.simulator.sim_container import Simulation\n", + " from primaite.simulator.network.hardware.nodes.computer import Computer\n", + " from primaite.game.actor.observations import FileObservation\n", + "\n", + " sim = Simulation()\n", + " pc = Computer(hostname=\"beep\", ip_address=\"123.123.123.123\", subnet_mask=\"255.255.255.0\")\n", + " sim.network.add_node(pc)\n", + " f = pc.file_system.create_file(file_name=\"dog.png\")\n", + "\n", + " dog_file_obs = FileObservation(where=['network','nodes',pc.uuid,'file_system'])\n", + " print(sim.describe_state())\n", + "test_file_observation()" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "ename": "AttributeError", + "evalue": "'NIC' object has no attribute 'gateway'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)", + "File \u001b[0;32m~/repos/PrimAITE/venv/lib/python3.10/site-packages/pydantic/main.py:718\u001b[0m, in \u001b[0;36mBaseModel.__getattr__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m 717\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 718\u001b[0m \u001b[39mreturn\u001b[39;00m pydantic_extra[item]\n\u001b[1;32m 719\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m \u001b[39mas\u001b[39;00m exc:\n", + "\u001b[0;31mKeyError\u001b[0m: 'gateway'", + "\nThe above exception was the direct cause of the following exception:\n", + "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", + "\u001b[1;32m/home/cade/repos/PrimAITE/test.ipynb Cell 6\u001b[0m line \u001b[0;36m1\n\u001b[1;32m 7\u001b[0m sim\u001b[39m.\u001b[39mnetwork\u001b[39m.\u001b[39madd_node(pc)\n\u001b[1;32m 8\u001b[0m f \u001b[39m=\u001b[39m pc\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mcreate_file(file_name\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mdog.png\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m---> 10\u001b[0m sim\u001b[39m.\u001b[39;49mdescribe_state()\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/sim_container.py:54\u001b[0m, in \u001b[0;36mSimulation.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 43\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 44\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 45\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 49\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 50\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 51\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 52\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 53\u001b[0m {\n\u001b[0;32m---> 54\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnetwork\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnetwork\u001b[39m.\u001b[39;49mdescribe_state(),\n\u001b[1;32m 55\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mdomain\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdomain\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 56\u001b[0m }\n\u001b[1;32m 57\u001b[0m )\n\u001b[1;32m 58\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/container.py:166\u001b[0m, in \u001b[0;36mNetwork.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 158\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 159\u001b[0m \u001b[39mProduce a dictionary describing the current state of the Network.\u001b[39;00m\n\u001b[1;32m 160\u001b[0m \n\u001b[1;32m 161\u001b[0m \u001b[39m:return: A dictionary capturing the current state of the Network and its child objects.\u001b[39;00m\n\u001b[1;32m 162\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 163\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 164\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 165\u001b[0m {\n\u001b[0;32m--> 166\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnodes\u001b[39m\u001b[39m\"\u001b[39m: {uuid:node\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, node \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 167\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mlinks\u001b[39m\u001b[39m\"\u001b[39m: {uuid:link\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, link \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlinks\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 168\u001b[0m }\n\u001b[1;32m 169\u001b[0m )\n\u001b[1;32m 170\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/container.py:166\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 158\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 159\u001b[0m \u001b[39mProduce a dictionary describing the current state of the Network.\u001b[39;00m\n\u001b[1;32m 160\u001b[0m \n\u001b[1;32m 161\u001b[0m \u001b[39m:return: A dictionary capturing the current state of the Network and its child objects.\u001b[39;00m\n\u001b[1;32m 162\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 163\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 164\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 165\u001b[0m {\n\u001b[0;32m--> 166\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnodes\u001b[39m\u001b[39m\"\u001b[39m: {uuid:node\u001b[39m.\u001b[39;49mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, node \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 167\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mlinks\u001b[39m\u001b[39m\"\u001b[39m: {uuid:link\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, link \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlinks\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 168\u001b[0m }\n\u001b[1;32m 169\u001b[0m )\n\u001b[1;32m 170\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:954\u001b[0m, in \u001b[0;36mNode.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 941\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 942\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 943\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 947\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 948\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 949\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 950\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 951\u001b[0m {\n\u001b[1;32m 952\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mhostname\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mhostname,\n\u001b[1;32m 953\u001b[0m \u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moperating_state\u001b[39m.\u001b[39mvalue,\n\u001b[0;32m--> 954\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mNICs\u001b[39m\u001b[39m\"\u001b[39m: {uuid: nic\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, nic \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnics\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 955\u001b[0m \u001b[39m# \"switch_ports\": {uuid, sp for uuid, sp in self.switch_ports.items()},\u001b[39;00m\n\u001b[1;32m 956\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mfile_system\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 957\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mapplications\u001b[39m\u001b[39m\"\u001b[39m: {uuid: app\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, app \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mapplications\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 958\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mservices\u001b[39m\u001b[39m\"\u001b[39m: {uuid: svc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, svc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 959\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mprocess\u001b[39m\u001b[39m\"\u001b[39m: {uuid: proc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, proc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprocesses\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 960\u001b[0m }\n\u001b[1;32m 961\u001b[0m )\n\u001b[1;32m 962\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:954\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 941\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 942\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 943\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 947\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 948\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 949\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 950\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 951\u001b[0m {\n\u001b[1;32m 952\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mhostname\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mhostname,\n\u001b[1;32m 953\u001b[0m \u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moperating_state\u001b[39m.\u001b[39mvalue,\n\u001b[0;32m--> 954\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mNICs\u001b[39m\u001b[39m\"\u001b[39m: {uuid: nic\u001b[39m.\u001b[39;49mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, nic \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnics\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 955\u001b[0m \u001b[39m# \"switch_ports\": {uuid, sp for uuid, sp in self.switch_ports.items()},\u001b[39;00m\n\u001b[1;32m 956\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mfile_system\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 957\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mapplications\u001b[39m\u001b[39m\"\u001b[39m: {uuid: app\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, app \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mapplications\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 958\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mservices\u001b[39m\u001b[39m\"\u001b[39m: {uuid: svc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, svc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 959\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mprocess\u001b[39m\u001b[39m\"\u001b[39m: {uuid: proc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, proc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprocesses\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 960\u001b[0m }\n\u001b[1;32m 961\u001b[0m )\n\u001b[1;32m 962\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:138\u001b[0m, in \u001b[0;36mNIC.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 125\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 126\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 127\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 131\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 132\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 133\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 134\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 135\u001b[0m {\n\u001b[1;32m 136\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mip_adress\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mip_address),\n\u001b[1;32m 137\u001b[0m \u001b[39m\"\u001b[39m\u001b[39msubnet_mask\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39msubnet_mask),\n\u001b[0;32m--> 138\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mgateway\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mgateway),\n\u001b[1;32m 139\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mmac_address\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmac_address,\n\u001b[1;32m 140\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mspeed\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mspeed,\n\u001b[1;32m 141\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mmtu\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmtu,\n\u001b[1;32m 142\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mwake_on_lan\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mwake_on_lan,\n\u001b[1;32m 143\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mdns_servers\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdns_servers,\n\u001b[1;32m 144\u001b[0m \u001b[39m\"\u001b[39m\u001b[39menabled\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39menabled,\n\u001b[1;32m 145\u001b[0m }\n\u001b[1;32m 146\u001b[0m )\n\u001b[1;32m 147\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", + "File \u001b[0;32m~/repos/PrimAITE/venv/lib/python3.10/site-packages/pydantic/main.py:720\u001b[0m, in \u001b[0;36mBaseModel.__getattr__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m 718\u001b[0m \u001b[39mreturn\u001b[39;00m pydantic_extra[item]\n\u001b[1;32m 719\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m \u001b[39mas\u001b[39;00m exc:\n\u001b[0;32m--> 720\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mAttributeError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m!r}\u001b[39;00m\u001b[39m object has no attribute \u001b[39m\u001b[39m{\u001b[39;00mitem\u001b[39m!r}\u001b[39;00m\u001b[39m'\u001b[39m) \u001b[39mfrom\u001b[39;00m \u001b[39mexc\u001b[39;00m\n\u001b[1;32m 721\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 722\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mhasattr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m, item):\n", + "\u001b[0;31mAttributeError\u001b[0m: 'NIC' object has no attribute 'gateway'" + ] + } + ], + "source": [ + "from primaite.simulator.sim_container import Simulation\n", + "from primaite.simulator.network.hardware.nodes.computer import Computer\n", + "from primaite.game.actor.observations import FileObservation\n", + "\n", + "sim = Simulation()\n", + "pc = Computer(hostname=\"beep\", ip_address=\"123.123.123.123\", subnet_mask=\"255.255.255.0\")\n", + "sim.network.add_node(pc)\n", + "f = pc.file_system.create_file(file_name=\"dog.png\")\n", + "\n", + "sim.describe_state()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "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.12" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/primaite/game/__init__.py b/src/primaite/game/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/game/actor/__init__.py b/src/primaite/game/actor/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/game/actor/actions.py b/src/primaite/game/actor/actions.py new file mode 100644 index 00000000..cefd9917 --- /dev/null +++ b/src/primaite/game/actor/actions.py @@ -0,0 +1,21 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from pydantic import BaseModel + + +class AbstractAction(BaseModel): + @abstractmethod + def __call__(self, action: Any) -> List[str]: + """_summary_ + + :param action: _description_ + :type action: Any + :return: _description_ + :rtype: List[str] + """ + ... + + +class ActionSpace: + ... diff --git a/src/primaite/game/actor/interface.py b/src/primaite/game/actor/interface.py new file mode 100644 index 00000000..1fe43a32 --- /dev/null +++ b/src/primaite/game/actor/interface.py @@ -0,0 +1,32 @@ +# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' +# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word +# If you disagree, make a comment in the PR review and we can discuss +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from pydantic import BaseModel + +from primaite.game.actor.actions import ActionSpace +from primaite.game.actor.observations import ObservationSpace +from primaite.game.actor.rewards import RewardFunction + + +class AbstractActor(BaseModel): + """Base class for scripted and RL agents.""" + + ... + + +class AbstractScriptedActor(AbstractActor): + """Base class for actors which generate their own behaviour.""" + + ... + + +class AbstractPuppetActor(AbstractActor): + """Base class for actors controlled via external messages, such as RL policies.""" + + ... + + +# class AbstractRLActor(AbstractPuppetActor): ?? diff --git a/src/primaite/game/actor/observations.py b/src/primaite/game/actor/observations.py new file mode 100644 index 00000000..bcde1c8d --- /dev/null +++ b/src/primaite/game/actor/observations.py @@ -0,0 +1,107 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, Hashable, List + +from pydantic import BaseModel + + +def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: + """ + Access an item from a deeply dictionary with a list of keys. + + For example, if the dictionary is {1: 'a', 2: {3: {4: 'b'}}}, then the key [2, 3, 4] would return 'b', and the key + [2, 3] would return {4: 'b'}. Raises a KeyError if specified key does not exist at any level of nesting. + + :param dictionary: Deeply nested dictionary + :type dictionary: Dict + :param keys: List of dict keys used to traverse the nested dict. Each item corresponds to one level of depth. + :type keys: List[Hashable] + :return: The value in the dictionary + :rtype: Any + """ + if not keys: + return dictionary + k = keys.pop(0) + try: + return access_from_nested_dict(dictionary[k], keys) + except (TypeError, KeyError): + raise KeyError(f"Cannot find requested key `{k}` in nested dictionary") + + +class AbstractObservation(BaseModel): + @abstractmethod + def __call__(self, state: Dict) -> Any: + """_summary_ + + :param state: _description_ + :type state: Dict + :return: _description_ + :rtype: Any + """ + ... + # receive state dict + + +class FileObservation(AbstractObservation): + where: List[str] + """Store information about where in the simulation state dictionary to find the relevatn information.""" + + def __call__(self, state: Dict) -> Dict: + file_state = access_from_nested_dict(state, self.where) + + +class ObservationSpace: + """Manage the observations of an Actor.""" + + ... + # what this class does: + # keep a list of observations + # create observations for an actor from the config + + +# Example YAML file for agent observation space +""" +arcd_gate: + rl_framework: SB3 + rl_algo: PPO + n_learn_steps: 128 + n_learn_episodes: 1000 + +game_layer: + agents: + - ref: client_1_green_user + type: GREEN + node_ref: client_1 + service: WebBrowser + pol: + - step: 1 + action: START + + - ref: client_1_data_manip_red_bot + node_ref: client_1 + service: DataManipulationBot + execution_definition: + - server_ip_address: 192.168.1.10 + - server_password: + - payload: 'ATTACK' + + pol: + - step: 75 + action: EXECUTE + + + + +simulation: + nodes: + - ref: client_1 + hostname: client_1 + node_type: Computer + ip_address: 192.168.10.100 + services: + - name: DataManipulationBot + links: + endpoint_a: + endpoint_b: 1524552-fgfg4147gdh-25gh4gd +rewards: + +""" diff --git a/src/primaite/game/actor/rewards.py b/src/primaite/game/actor/rewards.py new file mode 100644 index 00000000..1db54176 --- /dev/null +++ b/src/primaite/game/actor/rewards.py @@ -0,0 +1,20 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from pydantic import BaseModel + + +class AbstractReward(BaseModel): + def __call__(self, states: List[Dict]) -> float: + """_summary_ + + :param state: _description_ + :type state: Dict + :return: _description_ + :rtype: float + """ + ... + + +class RewardFunction(BaseModel): + ... diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py new file mode 100644 index 00000000..47ef4ce9 --- /dev/null +++ b/src/primaite/game/session.py @@ -0,0 +1,6 @@ +# What do? Be an entry point for using PrimAITE +# 1. parse monoconfig +# 2. craete simulation +# 3. create actors and configure their actions/observations/rewards/ anything else +# 4. Create connection with ARCD GATE +# 5. idk diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index b2037729..a6d2689b 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -3,6 +3,7 @@ from __future__ import annotations import math import os.path import shutil +from enum import Enum from pathlib import Path from typing import Dict, Optional @@ -42,6 +43,14 @@ def convert_size(size_bytes: int) -> str: return f"{s} {size_name[i]}" +class FileSystemItemHealthStatus(Enum): + GOOD = 1 + COMPROMISED = 2 + CORRUPT = 3 + RESTORING = 4 + REPAIRING = 5 + + class FileSystemItemABC(SimComponent): """ Abstract base class for file system items used in the file system simulation. @@ -51,6 +60,7 @@ class FileSystemItemABC(SimComponent): name: str "The name of the FileSystemItemABC." + health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD def describe_state(self) -> Dict: """ @@ -59,11 +69,7 @@ class FileSystemItemABC(SimComponent): :return: Current state of this object and child objects. """ state = super().describe_state() - state.update( - { - "name": self.name, - } - ) + state.update({"name": self.name, "health_status": self.health_status.value}) return state @property diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index c3a935b8..81ddf475 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -163,8 +163,8 @@ class Network(SimComponent): state = super().describe_state() state.update( { - "nodes": {i for i, node in self._node_id_map.items()}, - "links": {i: link.describe_state() for i, link in self._link_id_map.items()}, + "nodes": {uuid: node.describe_state() for uuid, node in self.nodes.items()}, + "links": {uuid: link.describe_state() for uuid, link in self.links.items()}, } ) return state From 53fd4ed82832761b8fd22e8407ddc6297da0b27a Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 21 Sep 2023 10:42:26 +0100 Subject: [PATCH 02/35] Finish File Observation --- sandbox.ipynb | 59 ++++++++++++++++--------- src/primaite/game/actor/observations.py | 15 +++++++ 2 files changed, 54 insertions(+), 20 deletions(-) diff --git a/sandbox.ipynb b/sandbox.ipynb index e7db5f4c..06e37664 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -66,26 +66,21 @@ "metadata": {}, "outputs": [ { - "ename": "AttributeError", - "evalue": "'NIC' object has no attribute 'gateway'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)", - "File \u001b[0;32m~/repos/PrimAITE/venv/lib/python3.10/site-packages/pydantic/main.py:718\u001b[0m, in \u001b[0;36mBaseModel.__getattr__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m 717\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 718\u001b[0m \u001b[39mreturn\u001b[39;00m pydantic_extra[item]\n\u001b[1;32m 719\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m \u001b[39mas\u001b[39;00m exc:\n", - "\u001b[0;31mKeyError\u001b[0m: 'gateway'", - "\nThe above exception was the direct cause of the following exception:\n", - "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", - "\u001b[1;32m/home/cade/repos/PrimAITE/test.ipynb Cell 6\u001b[0m line \u001b[0;36m1\n\u001b[1;32m 7\u001b[0m sim\u001b[39m.\u001b[39mnetwork\u001b[39m.\u001b[39madd_node(pc)\n\u001b[1;32m 8\u001b[0m f \u001b[39m=\u001b[39m pc\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mcreate_file(file_name\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mdog.png\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m---> 10\u001b[0m sim\u001b[39m.\u001b[39;49mdescribe_state()\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/sim_container.py:54\u001b[0m, in \u001b[0;36mSimulation.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 43\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 44\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 45\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 49\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 50\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 51\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 52\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 53\u001b[0m {\n\u001b[0;32m---> 54\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnetwork\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnetwork\u001b[39m.\u001b[39;49mdescribe_state(),\n\u001b[1;32m 55\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mdomain\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdomain\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 56\u001b[0m }\n\u001b[1;32m 57\u001b[0m )\n\u001b[1;32m 58\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/container.py:166\u001b[0m, in \u001b[0;36mNetwork.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 158\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 159\u001b[0m \u001b[39mProduce a dictionary describing the current state of the Network.\u001b[39;00m\n\u001b[1;32m 160\u001b[0m \n\u001b[1;32m 161\u001b[0m \u001b[39m:return: A dictionary capturing the current state of the Network and its child objects.\u001b[39;00m\n\u001b[1;32m 162\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 163\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 164\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 165\u001b[0m {\n\u001b[0;32m--> 166\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnodes\u001b[39m\u001b[39m\"\u001b[39m: {uuid:node\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, node \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 167\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mlinks\u001b[39m\u001b[39m\"\u001b[39m: {uuid:link\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, link \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlinks\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 168\u001b[0m }\n\u001b[1;32m 169\u001b[0m )\n\u001b[1;32m 170\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/container.py:166\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 158\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 159\u001b[0m \u001b[39mProduce a dictionary describing the current state of the Network.\u001b[39;00m\n\u001b[1;32m 160\u001b[0m \n\u001b[1;32m 161\u001b[0m \u001b[39m:return: A dictionary capturing the current state of the Network and its child objects.\u001b[39;00m\n\u001b[1;32m 162\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 163\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 164\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 165\u001b[0m {\n\u001b[0;32m--> 166\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mnodes\u001b[39m\u001b[39m\"\u001b[39m: {uuid:node\u001b[39m.\u001b[39;49mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, node \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 167\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mlinks\u001b[39m\u001b[39m\"\u001b[39m: {uuid:link\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, link \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlinks\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 168\u001b[0m }\n\u001b[1;32m 169\u001b[0m )\n\u001b[1;32m 170\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:954\u001b[0m, in \u001b[0;36mNode.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 941\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 942\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 943\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 947\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 948\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 949\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 950\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 951\u001b[0m {\n\u001b[1;32m 952\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mhostname\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mhostname,\n\u001b[1;32m 953\u001b[0m \u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moperating_state\u001b[39m.\u001b[39mvalue,\n\u001b[0;32m--> 954\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mNICs\u001b[39m\u001b[39m\"\u001b[39m: {uuid: nic\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, nic \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnics\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 955\u001b[0m \u001b[39m# \"switch_ports\": {uuid, sp for uuid, sp in self.switch_ports.items()},\u001b[39;00m\n\u001b[1;32m 956\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mfile_system\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 957\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mapplications\u001b[39m\u001b[39m\"\u001b[39m: {uuid: app\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, app \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mapplications\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 958\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mservices\u001b[39m\u001b[39m\"\u001b[39m: {uuid: svc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, svc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 959\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mprocess\u001b[39m\u001b[39m\"\u001b[39m: {uuid: proc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, proc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprocesses\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 960\u001b[0m }\n\u001b[1;32m 961\u001b[0m )\n\u001b[1;32m 962\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:954\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 941\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 942\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 943\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 947\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 948\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 949\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 950\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 951\u001b[0m {\n\u001b[1;32m 952\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mhostname\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mhostname,\n\u001b[1;32m 953\u001b[0m \u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moperating_state\u001b[39m.\u001b[39mvalue,\n\u001b[0;32m--> 954\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mNICs\u001b[39m\u001b[39m\"\u001b[39m: {uuid: nic\u001b[39m.\u001b[39;49mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, nic \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnics\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 955\u001b[0m \u001b[39m# \"switch_ports\": {uuid, sp for uuid, sp in self.switch_ports.items()},\u001b[39;00m\n\u001b[1;32m 956\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mfile_system\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfile_system\u001b[39m.\u001b[39mdescribe_state(),\n\u001b[1;32m 957\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mapplications\u001b[39m\u001b[39m\"\u001b[39m: {uuid: app\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, app \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mapplications\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 958\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mservices\u001b[39m\u001b[39m\"\u001b[39m: {uuid: svc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, svc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 959\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mprocess\u001b[39m\u001b[39m\"\u001b[39m: {uuid: proc\u001b[39m.\u001b[39mdescribe_state() \u001b[39mfor\u001b[39;00m uuid, proc \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprocesses\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 960\u001b[0m }\n\u001b[1;32m 961\u001b[0m )\n\u001b[1;32m 962\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/simulator/network/hardware/base.py:138\u001b[0m, in \u001b[0;36mNIC.describe_state\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 125\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 126\u001b[0m \u001b[39mProduce a dictionary describing the current state of this object.\u001b[39;00m\n\u001b[1;32m 127\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 131\u001b[0m \u001b[39m:rtype: Dict\u001b[39;00m\n\u001b[1;32m 132\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 133\u001b[0m state \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m()\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 134\u001b[0m state\u001b[39m.\u001b[39mupdate(\n\u001b[1;32m 135\u001b[0m {\n\u001b[1;32m 136\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mip_adress\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mip_address),\n\u001b[1;32m 137\u001b[0m \u001b[39m\"\u001b[39m\u001b[39msubnet_mask\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39msubnet_mask),\n\u001b[0;32m--> 138\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mgateway\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mstr\u001b[39m(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mgateway),\n\u001b[1;32m 139\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mmac_address\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmac_address,\n\u001b[1;32m 140\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mspeed\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mspeed,\n\u001b[1;32m 141\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mmtu\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmtu,\n\u001b[1;32m 142\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mwake_on_lan\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mwake_on_lan,\n\u001b[1;32m 143\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mdns_servers\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdns_servers,\n\u001b[1;32m 144\u001b[0m \u001b[39m\"\u001b[39m\u001b[39menabled\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39menabled,\n\u001b[1;32m 145\u001b[0m }\n\u001b[1;32m 146\u001b[0m )\n\u001b[1;32m 147\u001b[0m \u001b[39mreturn\u001b[39;00m state\n", - "File \u001b[0;32m~/repos/PrimAITE/venv/lib/python3.10/site-packages/pydantic/main.py:720\u001b[0m, in \u001b[0;36mBaseModel.__getattr__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m 718\u001b[0m \u001b[39mreturn\u001b[39;00m pydantic_extra[item]\n\u001b[1;32m 719\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m \u001b[39mas\u001b[39;00m exc:\n\u001b[0;32m--> 720\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mAttributeError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m!r}\u001b[39;00m\u001b[39m object has no attribute \u001b[39m\u001b[39m{\u001b[39;00mitem\u001b[39m!r}\u001b[39;00m\u001b[39m'\u001b[39m) \u001b[39mfrom\u001b[39;00m \u001b[39mexc\u001b[39;00m\n\u001b[1;32m 721\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 722\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mhasattr\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m, item):\n", - "\u001b[0;31mAttributeError\u001b[0m: 'NIC' object has no attribute 'gateway'" + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-09-21 10:41:35,339: Added node f03fec1b-927d-4d5a-8de9-1ef426052932 to Network f7400348-31e5-440e-8eb5-42366326d9d1\n" ] + }, + { + "data": { + "text/plain": [ + "{'health_status': 1}" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ @@ -98,7 +93,31 @@ "sim.network.add_node(pc)\n", "f = pc.file_system.create_file(file_name=\"dog.png\")\n", "\n", - "sim.describe_state()" + "state = sim.describe_state()\n", + "\n", + "dog_file_obs = FileObservation(where=['network','nodes',pc.uuid,'file_system', 'folders','root','files','dog.png'])\n", + "o = dog_file_obs(state)\n", + "o" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Dict(health_status:Discrete(6))" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "dog_file_obs.space" ] }, { diff --git a/src/primaite/game/actor/observations.py b/src/primaite/game/actor/observations.py index bcde1c8d..7303b07b 100644 --- a/src/primaite/game/actor/observations.py +++ b/src/primaite/game/actor/observations.py @@ -3,6 +3,8 @@ from typing import Any, Dict, Hashable, List from pydantic import BaseModel +from gym import spaces + def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: """ @@ -28,6 +30,7 @@ def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: class AbstractObservation(BaseModel): + @abstractmethod def __call__(self, state: Dict) -> Any: """_summary_ @@ -40,6 +43,12 @@ class AbstractObservation(BaseModel): ... # receive state dict + @property + @abstractmethod + def space(self) -> spaces.Space: + """Subclasses must define the shape that they expect""" + ... + class FileObservation(AbstractObservation): where: List[str] @@ -47,6 +56,12 @@ class FileObservation(AbstractObservation): def __call__(self, state: Dict) -> Dict: file_state = access_from_nested_dict(state, self.where) + observation = {'health_status':file_state['health_status']} + return observation + + @property + def space(self) -> spaces.Space: + return spaces.Dict({'health_status':spaces.Discrete(6)}) class ObservationSpace: From 9d4e41435d32e757c13990aebd480c82df71b56b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 25 Sep 2023 16:04:04 +0100 Subject: [PATCH 03/35] End-of-day commit --- example_config.yaml | 276 +++++++++++ src/primaite/game/actor/interface.py | 7 +- src/primaite/game/actor/observations.py | 456 +++++++++++++++--- .../simulator/network/hardware/base.py | 8 +- .../network/hardware/nodes/router.py | 21 +- .../simulator/system/services/service.py | 14 +- .../game_layer/test_observations.py | 20 + 7 files changed, 722 insertions(+), 80 deletions(-) create mode 100644 example_config.yaml create mode 100644 tests/integration_tests/game_layer/test_observations.py diff --git a/example_config.yaml b/example_config.yaml new file mode 100644 index 00000000..6c02031a --- /dev/null +++ b/example_config.yaml @@ -0,0 +1,276 @@ +training_config: + rl_framework: SB3 + rl_algo: PPO + n_learn_steps: 128 + n_learn_episodes: 1000 + +game_config: + ports: + - ARP + - DNS + - POSTGRES_SERVER + protocols: + - ICMP + - TCP + + agents: + - ref: client_1_green_user + team: GREEN + team: SCRIPTED_GREEN_ + observation_space: + ... + action_space: + ... + reward_function: + - type: null_reward + # node_ref: client_1 + # service: WebBrowser + # pol: + # - step: 1 + # action: START + + - ref: client_1_data_manipulation_red_bot + team: RED + type: SCRIPTED_RED_ + observation_space: + network: + nodes: + - ref: client_1 + - logon_status + - operating_status + services: + - ref: data_manipulation_bot + - operating_status + - health_status + folders: + files: {} + nics: {} + + action_space: + actions: + - DO_NOTHING + network: + nodes: + - ref: client_1 + actions: + - SCAN + - LOGON + - LOGOFF + services: + - ref: data_manipulation_bot + actions: + - type: COMPROMISE + execution_definition: + server_ip: 192.168.1.14 + payload: "DROP TABLE IF EXISTS user;" + success_rate: 80% + folders: + files: {} + reward_function: null + options: # options specific to this particular agent type, basically args of __init__(self) + start_step: 25 + frequency: 20 + variance: 5 + + + + + - ref: defender + team: blue + type: GATE_RL_AGENT + observation_space: + network: + nodes: + - ref: + action_space: + ... + reward_function: + ... + + + + + +simulation: + network: + nodes: + + - ref: router_1 + type: router + hostname: router_1 + num_ports: 5 + ports: + 1: + ip_address: 192.168.1.1 + subnet_mask: 255.255.255.0 + 2: + ip_address: 192.168.1.1 + subnet_mask: 255.255.255.0 + acl: + 0: + action: PERMIT + src_port: POSTGRES_SERVER + dst_port: POSTGRES_SERVER + 1: + action: PERMIT + src_port: DNS + dst_port: DNS + 22: + action: PERMIT + src_port: ARP + dst_port: ARP + 23: + action: PERMIT + protocol: ICMP + + - ref: switch_1 + type: swtich + hostname: switch_1 + num_ports: 8 + + - ref: switch_2 + type: switch + hostname: switch_2 + num_ports: 8 + + - ref: domain_controller + type: server + hostname: domain_controller + ip_address: 192.168.1.10 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + services: + - ref: domain_controller_dns_server + type: dns_server + options: + domain_mapping: + - arcd.com: 192.168.1.12 # web server + + + - ref: web_server + type: server + hostname: web_server + ip_address: 192.168.1.12 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.10 + dns_server: 192.168.1.10 + services: + - ref: web_server_database_client + type: database_client + options: + db_server_ip: 192.168.1.14 + + - ref: database_server + type: server + hostname: database_server + ip_address: 192.168.1.14 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + services: + - ref: database_service + type: database_service + + + - ref: backup_server + type: node + hostname: backup_server + ip_address: 192.168.1.16 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + services: + - ref: backup_service + type: database_backup + + - ref: security_suite + type: server + hostname: security_suite + ip_address: 192.168.1.110 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + nics: + 2: + ip_address: 192.168.10.110 + subnet_mask: 255.255.255.0 + + + - ref: client_1 + type: computer + hostname: client_1 + ip_address: 192.168.10.21. + subnet_mask: 255.255.255.0 + default_gateway: 192.168.10.1 + dns_server: 192.168.1.10 + services: + - ref: data_manipulation_bot + type: data_manipulation_bot + - ref: client_1_dns_client + type: dns_client + + - ref: client_2 + type: computer + hostname: client_2 + ip_address: 192.168.10.22 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.10.1 + dns_server: 192.168.1.10 + services: + - ref: web_browser + type: web_browser + - ref: client_2_dns_client + type: dns_client + + + links: + - ref: router_1___switch_1 + endpoint_a: router_1 + endpoint_a_port: 1 + endpoint_b: switch_1 + endpoint_b_port: 8 + - ref: router_1___switch_2 + endpoint_a: router_1 + endpoint_a_port: 2 + endpoint_b: switch_2 + endpoint_b_port: 8 + - ref: switch_1___domain_controller + endpoint_a: switch_1 + endpoint_a_port: 1 + endpoint_b: domain_controller + endpoint_b_port: 1 + - ref: switch_1___web_server + endpoint_a: switch_1 + endpoint_a_port: 2 + endpoint_b: web_server + endpoint_b_port: 1 + - ref: switch_1___database_server + endpoint_a: switch_1 + endpoint_a_port: 3 + endpoint_b: database_server + endpoint_b_port: 1 + - ref: switch_1___backup_server + endpoint_a: switch_1 + endpoint_a_port: 4 + endpoint_b: backup_server + endpoint_b_port: 1 + - ref: switch_1___security_suite + endpoint_a: switch_1 + endpoint_a_port: 7 + endpoint_b: security_suite + endpoint_b_port: 1 + - ref: switch_2___client_1 + endpoint_a: switch_2 + endpoint_a_port: 1 + endpoint_b: client_1 + endpoint_b_port: 1 + - ref: switch_2___client_2 + endpoint_a: switch_2 + endpoint_a_port: 2 + endpoint_b: client_2 + endpoint_b_port: 1 + - ref: switch_2___security_suite + endpoint_a: switch_2 + endpoint_a_port: 7 + endpoint_b: security_suite + endpoint_b_port: 2 diff --git a/src/primaite/game/actor/interface.py b/src/primaite/game/actor/interface.py index 1fe43a32..d1245e71 100644 --- a/src/primaite/game/actor/interface.py +++ b/src/primaite/game/actor/interface.py @@ -11,10 +11,13 @@ from primaite.game.actor.observations import ObservationSpace from primaite.game.actor.rewards import RewardFunction -class AbstractActor(BaseModel): +class AbstractActor(ABC): """Base class for scripted and RL agents.""" - ... + def __init__(self) -> None: + self.action_space = ActionSpace + self.observation_space = ObservationSpace + self.reward_function = RewardFunction class AbstractScriptedActor(AbstractActor): diff --git a/src/primaite/game/actor/observations.py b/src/primaite/game/actor/observations.py index 7303b07b..4d4796e1 100644 --- a/src/primaite/game/actor/observations.py +++ b/src/primaite/game/actor/observations.py @@ -1,9 +1,16 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, Hashable, List - -from pydantic import BaseModel +from typing import Any, Dict, Hashable, List, Optional from gym import spaces +from pydantic import BaseModel + +from primaite.simulator.sim_container import Simulation + +NOT_PRESENT_IN_STATE = object() +""" +Need an object to return when the sim state does not contain a requested value. Cannot use None because sometimes +the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is a sentinel for this purpose. +""" def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: @@ -20,19 +27,17 @@ def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: :return: The value in the dictionary :rtype: Any """ - if not keys: + if len(keys) == 0: return dictionary k = keys.pop(0) - try: - return access_from_nested_dict(dictionary[k], keys) - except (TypeError, KeyError): - raise KeyError(f"Cannot find requested key `{k}` in nested dictionary") + if k not in dictionary: + return NOT_PRESENT_IN_STATE + return access_from_nested_dict(dictionary[k], keys) -class AbstractObservation(BaseModel): - +class AbstractObservation(ABC): @abstractmethod - def __call__(self, state: Dict) -> Any: + def observe(self, state: Dict) -> Any: """_summary_ :param state: _description_ @@ -41,7 +46,6 @@ class AbstractObservation(BaseModel): :rtype: Any """ ... - # receive state dict @property @abstractmethod @@ -51,72 +55,396 @@ class AbstractObservation(BaseModel): class FileObservation(AbstractObservation): - where: List[str] - """Store information about where in the simulation state dictionary to find the relevatn information.""" + def __init__(self, where: List[str] = []) -> None: + """ + _summary_ - def __call__(self, state: Dict) -> Dict: + :param where: Store information about where in the simulation state dictionary to find the relevatn information. + Optional. If None, this corresponds that the file does not exist and the observation will be populated with + zeroes. + + A typical location for a file looks like this: + ['network','nodes',,'file_system', 'folders',,'files',] + :type where: Optional[List[str]] + """ + super().__init__() + self.where: List[str] = where + self.default_observation: spaces.Space = {"health_status": 0} + "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted." + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation file_state = access_from_nested_dict(state, self.where) - observation = {'health_status':file_state['health_status']} - return observation + if file_state is NOT_PRESENT_IN_STATE: + return self.default_observation + return {"health_status": file_state["health_status"]} @property def space(self) -> spaces.Space: - return spaces.Dict({'health_status':spaces.Discrete(6)}) + return spaces.Dict({"health_status": spaces.Discrete(6)}) + + +class ServiceObservation(AbstractObservation): + default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} + "Default observation is what should be returned when the service doesn't exist." + + def __init__(self, where: List[str] = []) -> None: + """ + :param where: Store information about where in the simulation state dictionary to find the relevant information. + Optional. If None, this corresponds that the file does not exist and the observation will be populated with + zeroes. + + A typical location for a service looks like this: + `['network','nodes',,'servics', ]` + :type where: Optional[List[str]] + """ + super().__init__() + self.where: List[str] = where + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + + service_state = access_from_nested_dict(state, self.where) + if service_state is NOT_PRESENT_IN_STATE: + return self.default_observation + return {"operating_status": service_state["operating_status"], "health_status": service_state["health_status"]} + + @property + def space(self) -> spaces.Space: + return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)}) + + +class LinkObservation(AbstractObservation): + default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} + "Default observation is what should be returned when the link doesn't exist." + + def __init__(self, where: List[str] = []) -> None: + """ + :param where: Store information about where in the simulation state dictionary to find the relevant information. + Optional. If None, this corresponds that the file does not exist and the observation will be populated with + zeroes. + + A typical location for a service looks like this: + `['network','nodes',,'servics', ]` + :type where: Optional[List[str]] + """ + super().__init__() + self.where: List[str] = where + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + + link_state = access_from_nested_dict(state, self.where) + if link_state is NOT_PRESENT_IN_STATE: + return self.default_observation + + bandwidth = link_state["bandwidth"] + load = link_state["current_load"] + utilisation_fraction = load / bandwidth + # 0 is UNUSED, 1 is 0%-10%. 2 is 10%-20%. 3 is 20%-30%. And so on... 10 is exactly 100% + utilisation_category = int(utilisation_fraction * 10) + 1 + + # TODO: once the links support separte load per protocol, this needs amendment to reflect that. + return {"protocols": {"all": {"load": utilisation_category}}} + + @property + def space(self) -> spaces.Space: + return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})}) + + +class FolderObservation(AbstractObservation): + def __init__(self, where: List[str] = [], files: List[FileObservation] = []) -> None: + """Initialise folder Observation, including files inside of the folder. + + :param where: Where in the simulation state dictionary to find the relevant information for this folder. + A typical location for a file looks like this: + ['network','nodes',,'file_system', 'folders',] + :type where: Optional[List[str]] + :param max_files: As size of the space must remain static, define max files that can be in this folder + , defaults to 5 + :type max_files: int, optional + :param file_positions: Defines the positioning within the observation space of particular files. This ensures + that even if new files are created, the existing files will always occupy the same space in the observation + space. The keys must be between 1 and max_files. Providing file_positions will reserve a spot in the + observation space for a file with that name, even if it's temporarily deleted, if it reappears with the same + name, it will take the position defined in this dict. Defaults to {} + :type file_positions: Dict[int, str], optional + """ + super().__init__() + + self.where: List[str] = where + + self.files: List[FileObservation] = files + + self.default_observation = { + "health_status": 0, + "FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)}, + } + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + folder_state = access_from_nested_dict(state, self.where) + if folder_state is NOT_PRESENT_IN_STATE: + return self.default_observation + + health_status = folder_state["health_status"] + + obs = {} + + obs["health_status"] = health_status + obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)} + + return obs + + @property + def space(self) -> spaces.Space: + return spaces.Dict( + { + "health_status": spaces.Discrete(6), + "FILES": spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)}), + } + ) + + +class NicObservation(AbstractObservation): + default_observation: spaces.Space = {"nic_status": 0} + + def __init__(self, where: List[str] = []) -> None: + super.__init__() + self.where: List[str] = where + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + nic_state = access_from_nested_dict(state, self.where) + if nic_state is NOT_PRESENT_IN_STATE: + return self.default_observation + else: + return {"nic_status": 1 if nic_state["enabled"] else 2} + + @property + def space(self) -> spaces.Space: + return spaces.Dict({"nic_status": spaces.Discrete(3)}) + + +class NodeObservation(AbstractObservation): + def __init__( + self, + where: List[str] = [], + services: List[ServiceObservation] = [], + folders: List[FolderObservation] = [], + nics: List[NicObservation] = [], + ) -> None: + """ + Configurable observation for a node in the simulation. + + :param where: Where in the simulation state dictionary for find relevant information for this observation. + A typical location for a node looks like this: + ['network','nodes',]. If empty list, a default null observation will be output, defaults to [] + :type where: List[str], optional + :param services: Mapping between position in observation space and service UUID, defaults to {} + :type services: Dict[int,str], optional + :param max_services: Max number of services that can be presented in observation space for this node, defaults to 2 + :type max_services: int, optional + :param folders: Mapping between position in observation space and folder name, defaults to {} + :type folders: Dict[int,str], optional + :param max_folders: Max number of folders in this node's obs space, defaults to 2 + :type max_folders: int, optional + :param nics: Mapping between position in observation space and NIC UUID, defaults to {} + :type nics: Dict[int,str], optional + :param max_nics: Max number of NICS in this node's obs space, defaults to 5 + :type max_nics: int, optional + """ + super.__init__() + self.where: List[str] = where + + self.services: List[ServiceObservation] = services + self.folders: List[FolderObservation] = folders + self.nics: List[NicObservation] = nics + + self.default_observation: Dict = { + "SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)}, + "FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)}, + "NICS": {i + 1: n.default_observation for i, n in enumerate(self.nics)}, + "operating_status": 0, + } + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + + node_state = access_from_nested_dict(state, self.where) + if node_state is NOT_PRESENT_IN_STATE: + return self.default_observation + + obs = {} + + obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)} + obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)} + obs["operating_status"] = node_state["operating_state"] + obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)} + + return obs + + @property + def space(self) -> spaces.Space: + return spaces.Dict( + { + "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}), + "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}), + "operating_status": spaces.Discrete(0), + "NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}), + } + ) + + +class AclObservation(AbstractObservation): + # TODO: should where be optional, and we can use where=None to pad the observation space? + # definitely the current approach does not support tracking files that aren't specified by name, for example + # if a file is created at runtime, we have currently got no way of telling the observation space to track it. + # this needs adding, but not for the MVP. + def __init__( + self, nodes: List[str], ports: List[int], protocols: list[str], where: List[str] = [], num_rules: int = 10 + ) -> None: + super().__init__() + self.where: List[str] = where + self.num_rules: int = num_rules + self.node_to_id: Dict[str, int] = {node: i + 1 for i, node in enumerate(nodes)} + "List of node IP addresses, order in this list determines how they are converted to an ID" + self.port_to_id: Dict[int, int] = {port: i + 1 for i, port in enumerate(ports)} + "List of ports which are part of the game that define the ordering when converting to an ID" + self.protocol_to_id: Dict[str, int] = {protocol: i + 1 for i, protocol in enumerate(protocols)} + "List of protocols which are part of the game, defines ordering when converting to an ID" + self.default_observation: spaces.Space = spaces.Dict( + { + "RULES": spaces.Dict( + { + i + + 1: spaces.Dict( + { + "position": i, + "permission": 0, + "source_node_id": 0, + "source_port": 0, + "dest_node_id": 0, + "dest_port": 0, + "protocol": 0, + } + ) + for i in range(self.num_rules) + } + ) + } + ) + + def observe(self, state: Dict) -> Dict: + if not self.where: + return self.default_observation + acl_state: Dict = access_from_nested_dict(state, self.where) + if acl_state is NOT_PRESENT_IN_STATE: + return self.default_observation + + obs = {} + obs["RULES"] = {} + for i, rule_state in acl_state.items(): + if rule_state is None: + obs["RULES"][i + 1] = { + "position": i, + "permission": 0, + "source_node_id": 0, + "source_port": 0, + "dest_node_id": 0, + "dest_port": 0, + "protocol": 0, + } + else: + obs["RULES"][i + 1] = { + "position": i, + "permission": rule_state["action"], + "source_node_id": self.node_to_id[rule_state["src_ip_address"]], + "source_port": self.port_to_id[rule_state["src_port"]], + "dest_node_id": self.node_to_id[rule_state["dst_ip_address"]], + "dest_port": self.port_to_id[rule_state["dst_port"]], + "protocol": self.protocol_to_id[rule_state["protocol"]], + } + return obs + + @property + def space(self) -> spaces.Space: + return spaces.Dict( + { + "RULE": spaces.Dict( + { + i + + 1: spaces.Dict( + { + "position": spaces.Discrete(self.num_rules), + "permission": spaces.Discrete(3), + "source_node_id": spaces.Discrete(len(self.nodes) + 1), + "source_port": spaces.Discrete(len(self.ports) + 1), + "dest_node_id": spaces.Discrete(len(self.nodes) + 1), + "dest_port": spaces.Discrete(len(self.ports) + 1), + "protocol": spaces.Discrete(len(self.protocols) + 1), + } + ) + for i in range(self.num_rules) + } + ) + } + ) + + +class ICSObservation(AbstractObservation): + def observe(self, state: Dict) -> Any: + return 0 + + @property + def space(self) -> spaces.Space: + return spaces.Discrete(1) class ObservationSpace: - """Manage the observations of an Actor.""" + """ + Manage the observations of an Actor. + + The observation space has the purpose of: + 1. Reading the outputted state from the PrimAITE Simulation. + 2. Selecting parts of the simulation state that are requested by the simulation config + 3. Formatting this information so an actor can use it to make decisions. + """ ... + # what this class does: # keep a list of observations # create observations for an actor from the config + def __init__( + self, + simulation: Simulation, + nodes: List[NodeObservation] = [], + links: List[LinkObservation] = [], + acl: Optional[AclObservation] = None, + ics: Optional[ICSObservation] = None, + ) -> None: + self.simulation: Simulation = simulation + self.parts: Dict[str, AbstractObservation] = {} + self.nodes: List[NodeObservation] = nodes + self.links: List[LinkObservation] = links + self.acl: Optional[AclObservation] = acl + self.ics: Optional[ICSObservation] = ics -# Example YAML file for agent observation space -""" -arcd_gate: - rl_framework: SB3 - rl_algo: PPO - n_learn_steps: 128 - n_learn_episodes: 1000 + def observe(self) -> None: + ... -game_layer: - agents: - - ref: client_1_green_user - type: GREEN - node_ref: client_1 - service: WebBrowser - pol: - - step: 1 - action: START + @property + def space(self) -> None: + ... - - ref: client_1_data_manip_red_bot - node_ref: client_1 - service: DataManipulationBot - execution_definition: - - server_ip_address: 192.168.1.10 - - server_password: - - payload: 'ATTACK' - - pol: - - step: 75 - action: EXECUTE - - - - -simulation: - nodes: - - ref: client_1 - hostname: client_1 - node_type: Computer - ip_address: 192.168.10.100 - services: - - name: DataManipulationBot - links: - endpoint_a: - endpoint_b: 1524552-fgfg4147gdh-25gh4gd -rewards: - -""" + @classmethod + def from_config(self) -> None: + ... diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index a14d6a6d..f5ba8444 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -855,14 +855,14 @@ class ICMP: class NodeOperatingState(Enum): """Enumeration of Node Operating States.""" - OFF = 0 - "The node is powered off." ON = 1 "The node is powered on." - SHUTTING_DOWN = 2 - "The node is in the process of shutting down." + OFF = 2 + "The node is powered off." BOOTING = 3 "The node is in the process of booting up." + SHUTTING_DOWN = 4 + "The node is in the process of shutting down." class Node(SimComponent): diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 53b9b176..7870caab 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -58,7 +58,14 @@ class ACLRule(SimComponent): :return: A dictionary representing the current state. """ - pass + state = super().describe_state() + state["action"] = self.action.value + state["protocol"] = self.protocol.value + state["src_ip_address"] = self.src_ip_address + state["src_port"] = self.src_port.value + state["dst_ip_address"] = self.dst_ip_address + state["dst_port"] = self.dst_port.value + return state class AccessControlList(SimComponent): @@ -123,7 +130,12 @@ class AccessControlList(SimComponent): :return: A dictionary representing the current state. """ - pass + state = super().describe_state() + state["implicit_action"] = self.implicit_action.value + state["implicit_rule"] = self.implicit_rule.describe_state() + state["max_acl_rules"] = self.max_acl_rules + state["acl"] = {i: r.describe_state() if isinstance(r, ACLRule) else None for i, r in enumerate(self._acl)} + return state @property def acl(self) -> List[Optional[ACLRule]]: @@ -648,7 +660,10 @@ class Router(Node): :return: A dictionary representing the current state. """ - pass + state = super().describe_state() + state["num_ports"] = (self.num_ports,) + state["acl"] = (self.acl.describe_state(),) + return state def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: """ diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 20b92027..fb12fc3d 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -15,14 +15,14 @@ class ServiceOperatingState(Enum): "The service is currently running." STOPPED = 2 "The service is not running." - INSTALLING = 3 - "The service is being installed or updated." - RESTARTING = 4 - "The service is in the process of restarting." - PAUSED = 5 + PAUSED = 3 "The service is temporarily paused." - DISABLED = 6 + DISABLED = 4 "The service is disabled and cannot be started." + INSTALLING = 5 + "The service is being installed or updated." + RESTARTING = 6 + "The service is in the process of restarting." class Service(IOSoftware): @@ -60,7 +60,7 @@ class Service(IOSoftware): :rtype: Dict """ state = super().describe_state() - state.update({"operating_state": self.operating_state.name}) + state.update({"operating_state": self.operating_state.value}) return state def reset_component_for_episode(self, episode: int): diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py new file mode 100644 index 00000000..4f778f78 --- /dev/null +++ b/tests/integration_tests/game_layer/test_observations.py @@ -0,0 +1,20 @@ +from gym import spaces + +from primaite.game.actor.observations import FileObservation +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.sim_container import Simulation + + +def test_file_observation(): + sim = Simulation() + pc = Computer(hostname="beep", ip_address="123.123.123.123", subnet_mask="255.255.255.0") + sim.network.add_node(pc) + f = pc.file_system.create_file(file_name="dog.png") + + state = sim.describe_state() + + dog_file_obs = FileObservation( + where=["network", "nodes", pc.uuid, "file_system", "folders", "root", "files", "dog.png"] + ) + assert dog_file_obs(state) == {"health_status": 1} + assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)}) From 493014ca19beec127689aaa9822a2b3719b4905e Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 25 Sep 2023 17:57:47 +0100 Subject: [PATCH 04/35] draft yaml parser --- example_config.yaml | 101 +++++++++++------- sandbox.ipynb | 250 +++++++++++++++++++++++++++++++++++++------- 2 files changed, 279 insertions(+), 72 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 6c02031a..f0957718 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -17,17 +17,26 @@ game_config: - ref: client_1_green_user team: GREEN team: SCRIPTED_GREEN_ - observation_space: - ... + observation_space: null action_space: - ... - reward_function: - - type: null_reward - # node_ref: client_1 - # service: WebBrowser - # pol: - # - step: 1 - # action: START + actions: + - type: DONOTHING + nodes: + - ref: client_2 + actions: + - type: LOGON + - type: LOGOFF + applications: + - ref: client_2_web_browser + actions: + - type: EXECUTE + execution_definition: + target_address: arcd.com + reward_function: null + agent_settings: + start_step: 5 + frequency: 4 + variance: 3 - ref: client_1_data_manipulation_red_bot team: RED @@ -36,38 +45,37 @@ game_config: network: nodes: - ref: client_1 + observations: - logon_status - operating_status services: - ref: data_manipulation_bot + observations: - operating_status - health_status - folders: - files: {} - nics: {} - + folders: {} action_space: actions: - - DO_NOTHING + - type: DO_NOTHING network: nodes: - - ref: client_1 + - ref: client_1 + actions: + - type: SCAN + - type: LOGON + - type: LOGOFF + services: + - ref: data_manipulation_bot actions: - - SCAN - - LOGON - - LOGOFF - services: - - ref: data_manipulation_bot - actions: - - type: COMPROMISE - execution_definition: - server_ip: 192.168.1.14 - payload: "DROP TABLE IF EXISTS user;" - success_rate: 80% - folders: - files: {} + - type: COMPROMISE + execution_definition: + server_ip: 192.168.1.14 + payload: "DROP TABLE IF EXISTS user;" + success_rate: 80% + folders: + files: {} reward_function: null - options: # options specific to this particular agent type, basically args of __init__(self) + agent_settings: # options specific to this particular agent type, basically args of __init__(self) start_step: 25 frequency: 20 variance: 5 @@ -81,11 +89,32 @@ game_config: observation_space: network: nodes: - - ref: + - ref: router_1 #TODO: more sub-options here + - ref: switch_1 + - ref: switch_2 + - ref: domain_controller + - ref: web_server + - ref: database_server + - ref: backup_server + - ref: security_suite + - ref: client_1 + - ref: client_2 + links: + - ref: ... # + acl: ... # + ics: ... # + + action_space: - ... + actions: + - type: DO_NOTHING + network: + nodes: + - ref: router_1 reward_function: - ... + # ... + agent_settings: + # ... @@ -173,7 +202,7 @@ simulation: - ref: backup_server - type: node + type: server hostname: backup_server ip_address: 192.168.1.16 subnet_mask: 255.255.255.0 @@ -199,7 +228,7 @@ simulation: - ref: client_1 type: computer hostname: client_1 - ip_address: 192.168.10.21. + ip_address: 192.168.10.21 subnet_mask: 255.255.255.0 default_gateway: 192.168.10.1 dns_server: 192.168.1.10 @@ -217,7 +246,7 @@ simulation: default_gateway: 192.168.10.1 dns_server: 192.168.1.10 services: - - ref: web_browser + - ref: client_2_web_browser type: web_browser - ref: client_2_dns_client type: dns_client diff --git a/sandbox.ipynb b/sandbox.ipynb index 06e37664..91edb829 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -1,14 +1,31 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import yaml" + ] + }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "from primaite.simulator.network.networks import arcd_uc2_network\n", - "%load_ext autoreload\n", - "%autoreload 2" + "from primaite.simulator.network.networks import arcd_uc2_network\n" ] }, { @@ -62,31 +79,13 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-09-21 10:41:35,339: Added node f03fec1b-927d-4d5a-8de9-1ef426052932 to Network f7400348-31e5-440e-8eb5-42366326d9d1\n" - ] - }, - { - "data": { - "text/plain": [ - "{'health_status': 1}" - ] - }, - "execution_count": 1, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "from primaite.simulator.sim_container import Simulation\n", "from primaite.simulator.network.hardware.nodes.computer import Computer\n", - "from primaite.game.actor.observations import FileObservation\n", + "from primaite.game.actor.observations import FileObservation, FolderObservation\n", "\n", "sim = Simulation()\n", "pc = Computer(hostname=\"beep\", ip_address=\"123.123.123.123\", subnet_mask=\"255.255.255.0\")\n", @@ -96,28 +95,207 @@ "state = sim.describe_state()\n", "\n", "dog_file_obs = FileObservation(where=['network','nodes',pc.uuid,'file_system', 'folders','root','files','dog.png'])\n", - "o = dog_file_obs(state)\n", - "o" + "root_folder_obs = FolderObservation(where=['network','nodes',pc.uuid,'file_system', 'folders','root'],files=[dog_file_obs])\n", + "print(dog_file_obs(state))\n", + "print(root_folder_obs(state))" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dog_file_obs.space" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "root_folder_obs.space" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "state" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import yaml" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open('example_config.yaml', 'r') as file:\n", + " conf = yaml.safe_load(file)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "conf['simulation']" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "import yaml\n", + "from primaite.simulator.sim_container import Simulation\n", + "from primaite.simulator.network.hardware.nodes.computer import Computer\n", + "from primaite.simulator.network.hardware.nodes.server import Server\n", + "from primaite.simulator.network.hardware.nodes.switch import Switch\n", + "from primaite.simulator.network.hardware.nodes.router import Router\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 33, "metadata": {}, "outputs": [ { - "data": { - "text/plain": [ - "Dict(health_status:Discrete(6))" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-09-25 17:38:39,385: Added node b5486651-1c6f-449a-8019-6a3641cfb998 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,391: Added node 1533c2f7-389e-4e03-95b3-9cf059086490 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,395: Added node 6b6c3b24-61d4-46ac-9364-11d726e50ccb to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,398: Added node a0bee8d0-2ab8-4e29-9a2c-23c6757b240c to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,401: Added node 7cb2c102-62ba-4859-94f2-5d724de38733 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,403: Added node bec38db7-520e-4044-93db-08308278d66f to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", + "2023-09-25 17:38:39,407: Added node ae0c2253-3ec8-48c3-b5d2-0b37c19c885d to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3\n", + "3\n", + "3\n" + ] } ], "source": [ - "dog_file_obs.space" + "# import yaml\n", + "\n", + "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", + "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", + "from primaite.simulator.network.transmission.transport_layer import Port\n", + "\n", + "\n", + "class PrimaiteSession:\n", + "\n", + " def __init__(self):\n", + " self.simulation: Simulation\n", + " self.agents = []\n", + "\n", + " @classmethod\n", + " def from_config(cls, cfg_path):\n", + " ref_to_uuid = {}\n", + "\n", + " game = cls()\n", + " with open(cfg_path, 'r') as file:\n", + " conf = yaml.safe_load(file)\n", + " \n", + " #1. create nodes \n", + " sim = Simulation()\n", + " net = sim.network\n", + " nodes_cfg = conf['simulation']['network']['nodes']\n", + " links_cfg = conf['simulation']['network']['links']\n", + " for node_cfg in nodes_cfg:\n", + " ref = node_cfg['ref']\n", + " n_type = node_cfg['type']\n", + " if n_type == 'computer':\n", + " new_node = Computer(hostname = node_cfg['hostname'], \n", + " ip_address = node_cfg['ip_address'], \n", + " subnet_mask = node_cfg['subnet_mask'], \n", + " default_gateway = node_cfg['default_gateway'],\n", + " dns_server = node_cfg['dns_server'])\n", + " elif n_type == 'server':\n", + " new_node = Server(hostname = node_cfg['hostname'], \n", + " ip_address = node_cfg['ip_address'], \n", + " subnet_mask = node_cfg['subnet_mask'], \n", + " default_gateway = node_cfg['default_gateway'],\n", + " dns_server = node_cfg.get('dns_server'))\n", + " elif n_type == 'switch':\n", + " new_node = Switch(hostname = node_cfg['hostname'],\n", + " num_ports = node_cfg.get('num_ports'))\n", + " elif n_type == 'router':\n", + " new_node = Router(hostname=node_cfg['hostname'],\n", + " num_ports = node_cfg.get('num_ports'))\n", + " if 'ports' in node_cfg:\n", + " for port_num, port_cfg in node_cfg['ports'].items():\n", + " new_node.configure_port(port=port_num, \n", + " ip_address=port_cfg['ip_address'],\n", + " subnet_mask=port_cfg['subnet_mask'])\n", + " if 'acl' in node_cfg:\n", + " for r_num, r_cfg in node_cfg['acl'].items():\n", + " new_node.acl.add_rule(\n", + " action = ACLAction[r_cfg['action']],\n", + " src_port = Port[r_cfg.get('port')],\n", + " dst_port = Port[r_cfg.get('port')],\n", + " protocol = IPProtocol[r_cfg.get('protocol')],\n", + " src_ip = r_cfg.get('ip_address'),\n", + " dst_ip = r_cfg.get('ip_address'),\n", + " position = r_num\n", + " )\n", + "\n", + "\n", + " try:\n", + " net.add_node(new_node)\n", + " ref_to_uuid[ref] = new_node.uuid\n", + " except BaseException:\n", + " print(3)\n", + "\n", + "\n", + " #2. start/setup simulation objects\n", + " #3. create agents\n", + " #4. set up agents' actions and observation spaces.\n", + " game.simulation = sim\n", + " return game\n", + "\n", + "s = PrimaiteSession.from_config('example_config.yaml')\n", + "# print(s.simulation.describe_state())" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'uuid': 'ceeb8791-b140-43d0-b59e-c3c3f533309b', 'network': {'uuid': 'ff176601-4e1d-4f89-8db4-33c0598ee105', 'nodes': {'6b9afe70-913b-40ce-9cee-1ee3648e43ce': {'uuid': '6b9afe70-913b-40ce-9cee-1ee3648e43ce', 'hostname': 'client_1', 'operating_state': 2, 'NICs': {'108c797d-32ca-4e93-8476-6b13cda6cf37': {'uuid': '108c797d-32ca-4e93-8476-6b13cda6cf37', 'ip_adress': '192.168.10.21', 'subnet_mask': '255.255.255.0', 'mac_address': 'af:5f:0c:00:d3:63', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '32062959-b2ed-4d24-b5a9-7e99b7ebfcfe', 'folders': {'root': {'uuid': '8876d59b-d46d-414d-9ae2-5e948f65b175', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '227d1fb7-fc64-4273-9817-0f32280a0859': {'uuid': '227d1fb7-fc64-4273-9817-0f32280a0859', 'hostname': 'client_2', 'operating_state': 2, 'NICs': {'22119571-b47d-4ffb-998c-62173c670f78': {'uuid': '22119571-b47d-4ffb-998c-62173c670f78', 'ip_adress': '192.168.10.22', 'subnet_mask': '255.255.255.0', 'mac_address': '7c:fe:81:20:96:96', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '1a6eb561-c7fc-40f0-a288-d56af08c8f0c', 'folders': {'root': {'uuid': 'd129d4a6-5098-41ee-b9b3-033895a2288c', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}}, 'links': {}}, 'domain': {'uuid': 'db0e6d12-7cc6-4828-ba9b-4110e7f14bc2', 'accounts': {}}}\n" + ] + } + ], + "source": [ + "print(s.simulation.describe_state())" ] }, { From f79ed99bd253335d59c82b81a3c48986ecedbdb4 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 25 Sep 2023 19:17:57 +0100 Subject: [PATCH 05/35] end of day --- sandbox.ipynb | 78 +++++++++++-------- .../network/hardware/nodes/router.py | 10 +-- 2 files changed, 51 insertions(+), 37 deletions(-) diff --git a/sandbox.ipynb b/sandbox.ipynb index 91edb829..96d12bae 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -157,7 +157,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 44, "metadata": {}, "outputs": [], "source": [ @@ -167,34 +167,33 @@ "from primaite.simulator.network.hardware.nodes.server import Server\n", "from primaite.simulator.network.hardware.nodes.switch import Switch\n", "from primaite.simulator.network.hardware.nodes.router import Router\n", - "\n" + "\n", + "from primaite.simulator.system.applications.database_client import DatabaseClient\n", + "from primaite.simulator.system.services.database_service import DatabaseService\n", + "from primaite.simulator.system.services.dns_client import DNSClient\n", + "from primaite.simulator.system.services.dns_server import DNSServer\n", + "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot\n" ] }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 42, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-25 17:38:39,385: Added node b5486651-1c6f-449a-8019-6a3641cfb998 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,391: Added node 1533c2f7-389e-4e03-95b3-9cf059086490 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,395: Added node 6b6c3b24-61d4-46ac-9364-11d726e50ccb to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,398: Added node a0bee8d0-2ab8-4e29-9a2c-23c6757b240c to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,401: Added node 7cb2c102-62ba-4859-94f2-5d724de38733 to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,403: Added node bec38db7-520e-4044-93db-08308278d66f to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n", - "2023-09-25 17:38:39,407: Added node ae0c2253-3ec8-48c3-b5d2-0b37c19c885d to Network 7c6e4724-1653-4db7-9bd0-44e8e380f1a1\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "3\n", - "3\n", - "3\n" + "2023-09-25 19:10:46,253: Added node 3d356af7-15f8-41d4-bb2a-423d5a2e5978 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,254::WARNING::primaite.simulator.network.container::181::Can't add node 3d356af7-15f8-41d4-bb2a-423d5a2e5978. It is already in the network.\n", + "2023-09-25 19:10:46,258: Added node 8a94447a-ccb3-47b2-b7ba-488e631f8246 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,262: Added node 4e3de72a-60ad-416f-abb7-da352a59d13b to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,264: Added node 21ae4eed-489f-43a7-af0a-f70000714c79 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,267: Added node c4b7d388-c13e-4f4e-be61-5cf05989da71 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,270: Added node 12dc9e72-fb7c-4a8f-ab79-d7574863fa16 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,273: Added node 37ff728f-64b5-4203-8177-cf27f51dc7c9 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,277: Added node fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a to Network 44950e57-81b0-4964-9b12-223592c785aa\n", + "2023-09-25 19:10:46,281: Added node dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6 to Network 44950e57-81b0-4964-9b12-223592c785aa\n" ] } ], @@ -204,6 +203,7 @@ "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", "from primaite.simulator.network.transmission.transport_layer import Port\n", + "from primaite.simulator.system.services.dns_server import DNSServer\n", "\n", "\n", "class PrimaiteSession:\n", @@ -226,7 +226,7 @@ " nodes_cfg = conf['simulation']['network']['nodes']\n", " links_cfg = conf['simulation']['network']['links']\n", " for node_cfg in nodes_cfg:\n", - " ref = node_cfg['ref']\n", + " node_ref = node_cfg['ref']\n", " n_type = node_cfg['type']\n", " if n_type == 'computer':\n", " new_node = Computer(hostname = node_cfg['hostname'], \n", @@ -253,22 +253,36 @@ " subnet_mask=port_cfg['subnet_mask'])\n", " if 'acl' in node_cfg:\n", " for r_num, r_cfg in node_cfg['acl'].items():\n", + " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, so that we can do\n", + " # both of these things once: check if a key is defined, access and convert it to a \n", + " # Port/IPProtocol. TODO Refactor\n", " new_node.acl.add_rule(\n", " action = ACLAction[r_cfg['action']],\n", - " src_port = Port[r_cfg.get('port')],\n", - " dst_port = Port[r_cfg.get('port')],\n", - " protocol = IPProtocol[r_cfg.get('protocol')],\n", - " src_ip = r_cfg.get('ip_address'),\n", - " dst_ip = r_cfg.get('ip_address'),\n", + " src_port = None if not (p:=r_cfg.get('src_port')) else Port[p],\n", + " dst_port = None if not (p:=r_cfg.get('dst_port')) else Port[p],\n", + " protocol = None if not (p:=r_cfg.get('protocol')) else IPProtocol[p],\n", + " src_ip_address = r_cfg.get('ip_address'),\n", + " dst_ip_address = r_cfg.get('ip_address'),\n", " position = r_num\n", " )\n", + " if 'services' in node_cfg:\n", + " for service_cfg in node_cfg['services']:\n", + " service_ref = service_cfg['ref']\n", + " service_type = service_cfg['type']\n", + " service_types_mapping = {\n", + " 'dns_server' : DNSServer,\n", + " 'database_client': DatabaseClient,\n", + " 'database_service': DatabaseService,\n", + " # 'database_backup': ,\n", + " 'data_manipulation_bot': DataManipulationBot,\n", + " 'dns_client': DNSClient,\n", + " # 'web_browser'\n", + " }\n", + " if service_type\n", "\n", "\n", - " try:\n", - " net.add_node(new_node)\n", - " ref_to_uuid[ref] = new_node.uuid\n", - " except BaseException:\n", - " print(3)\n", + " net.add_node(new_node)\n", + " ref_to_uuid[node_ref] = new_node.uuid\n", "\n", "\n", " #2. start/setup simulation objects\n", @@ -283,14 +297,14 @@ }, { "cell_type": "code", - "execution_count": 31, + "execution_count": 43, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "{'uuid': 'ceeb8791-b140-43d0-b59e-c3c3f533309b', 'network': {'uuid': 'ff176601-4e1d-4f89-8db4-33c0598ee105', 'nodes': {'6b9afe70-913b-40ce-9cee-1ee3648e43ce': {'uuid': '6b9afe70-913b-40ce-9cee-1ee3648e43ce', 'hostname': 'client_1', 'operating_state': 2, 'NICs': {'108c797d-32ca-4e93-8476-6b13cda6cf37': {'uuid': '108c797d-32ca-4e93-8476-6b13cda6cf37', 'ip_adress': '192.168.10.21', 'subnet_mask': '255.255.255.0', 'mac_address': 'af:5f:0c:00:d3:63', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '32062959-b2ed-4d24-b5a9-7e99b7ebfcfe', 'folders': {'root': {'uuid': '8876d59b-d46d-414d-9ae2-5e948f65b175', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '227d1fb7-fc64-4273-9817-0f32280a0859': {'uuid': '227d1fb7-fc64-4273-9817-0f32280a0859', 'hostname': 'client_2', 'operating_state': 2, 'NICs': {'22119571-b47d-4ffb-998c-62173c670f78': {'uuid': '22119571-b47d-4ffb-998c-62173c670f78', 'ip_adress': '192.168.10.22', 'subnet_mask': '255.255.255.0', 'mac_address': '7c:fe:81:20:96:96', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '1a6eb561-c7fc-40f0-a288-d56af08c8f0c', 'folders': {'root': {'uuid': 'd129d4a6-5098-41ee-b9b3-033895a2288c', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}}, 'links': {}}, 'domain': {'uuid': 'db0e6d12-7cc6-4828-ba9b-4110e7f14bc2', 'accounts': {}}}\n" + "{'uuid': '532b534a-78b7-4630-92f9-626d775ee4e2', 'network': {'uuid': '44950e57-81b0-4964-9b12-223592c785aa', 'nodes': {'3d356af7-15f8-41d4-bb2a-423d5a2e5978': {'uuid': '3d356af7-15f8-41d4-bb2a-423d5a2e5978', 'hostname': 'router_1', 'operating_state': 2, 'NICs': {'c52a1516-3488-4011-95ae-44e77935720d': {'uuid': 'c52a1516-3488-4011-95ae-44e77935720d', 'ip_adress': '192.168.1.1', 'subnet_mask': '255.255.255.0', 'mac_address': '12:75:33:08:ae:04', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '075a3387-b3c5-4e30-b7d5-f39b82b1c9d8': {'uuid': '075a3387-b3c5-4e30-b7d5-f39b82b1c9d8', 'ip_adress': '192.168.1.1', 'subnet_mask': '255.255.255.0', 'mac_address': 'ce:67:79:4a:48:90', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '3fe76cb7-84fe-4592-b527-be69c06c8064': {'uuid': '3fe76cb7-84fe-4592-b527-be69c06c8064', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '83:13:43:3a:e6:de', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, 'b4ad54f8-8508-4d34-9117-8f194a1076b1': {'uuid': 'b4ad54f8-8508-4d34-9117-8f194a1076b1', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '4f:1a:3f:76:c7:2c', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '21efd6b1-8b49-4b56-aafe-a3cbdf61e0fb': {'uuid': '21efd6b1-8b49-4b56-aafe-a3cbdf61e0fb', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '76:4d:65:e4:0d:11', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'df3b627a-76c3-4e60-83ef-81442fbf4643', 'folders': {'root': {'uuid': 'b21a0c20-e852-4d05-a0fe-9158f76f3a98', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}, 'num_ports': (5,), 'acl': ({'uuid': '2fc154d6-9ccf-4c60-b6d3-ac3fcce68fc8', 'implicit_action': 0, 'implicit_rule': {'uuid': '6a5553d0-37a6-4ac5-b5da-38f7926ffe7e', 'action': 0, 'protocol': None, 'src_ip_address': None, 'src_port': None, 'dst_ip_address': None, 'dst_port': None}, 'max_acl_rules': 25, 'acl': {0: {'uuid': '8204a738-6010-44a8-8df3-317d4a03f600', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 5432, 'dst_ip_address': None, 'dst_port': 5432}, 1: {'uuid': 'af7b39a5-e0ed-4047-b684-0cfbdc14f70b', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 53, 'dst_ip_address': None, 'dst_port': 53}, 2: None, 3: None, 4: None, 5: None, 6: None, 7: None, 8: None, 9: None, 10: None, 11: None, 12: None, 13: None, 14: None, 15: None, 16: None, 17: None, 18: None, 19: None, 20: None, 21: None, 22: {'uuid': '4c6320f5-d9f4-4958-9572-d2ebc06ee120', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 219, 'dst_ip_address': None, 'dst_port': 219}, 23: {'uuid': 'f39796b9-7511-4f23-b8a6-cbf0e87ec51a', 'action': 1, 'protocol': 'icmp', 'src_ip_address': None, 'src_port': None, 'dst_ip_address': None, 'dst_port': None}}},)}, '8a94447a-ccb3-47b2-b7ba-488e631f8246': {'uuid': '8a94447a-ccb3-47b2-b7ba-488e631f8246', 'num_ports': 8, 'ports': {1: {'uuid': '32eb0312-124d-42df-93bc-9cc25eb4b81f', 'mac_address': '1a:4e:83:7a:d7:4b', 'speed': 100, 'mtu': 1500, 'enabled': False}, 2: {'uuid': 'afe7a1b8-bfa6-42a2-b4aa-28187c67f7d3', 'mac_address': '9e:09:af:37:1d:e3', 'speed': 100, 'mtu': 1500, 'enabled': False}, 3: {'uuid': '98f46066-7e05-45cd-920e-54bbb3513fec', 'mac_address': '00:30:02:18:11:52', 'speed': 100, 'mtu': 1500, 'enabled': False}, 4: {'uuid': '65fb27f3-22ca-4bd3-898e-a2fdf49ec3b0', 'mac_address': 'f2:3f:0b:09:50:e1', 'speed': 100, 'mtu': 1500, 'enabled': False}, 5: {'uuid': 'f69c7045-40f7-490f-9c7c-280d0eb4fcd0', 'mac_address': '38:56:f5:dd:0a:a4', 'speed': 100, 'mtu': 1500, 'enabled': False}, 6: {'uuid': '6dbbac83-cd01-45ec-9ad6-352e1afee231', 'mac_address': 'aa:bb:c3:92:05:88', 'speed': 100, 'mtu': 1500, 'enabled': False}, 7: {'uuid': 'a39591ea-d50d-4649-ba90-f5fddd55b61d', 'mac_address': '60:67:46:45:21:cc', 'speed': 100, 'mtu': 1500, 'enabled': False}, 8: {'uuid': '488cdddb-9e5e-4fb1-8e03-cfcfd47b0cb3', 'mac_address': '0f:a0:4f:52:42:5c', 'speed': 100, 'mtu': 1500, 'enabled': False}}, 'mac_address_table': {}}, '4e3de72a-60ad-416f-abb7-da352a59d13b': {'uuid': '4e3de72a-60ad-416f-abb7-da352a59d13b', 'hostname': 'domain_controller', 'operating_state': 2, 'NICs': {'34962c55-cea2-479b-ac1f-61a683c82eb1': {'uuid': '34962c55-cea2-479b-ac1f-61a683c82eb1', 'ip_adress': '192.168.1.10', 'subnet_mask': '255.255.255.0', 'mac_address': '41:a0:d4:22:b9:81', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e3f270c9-2941-4592-892c-5662bbc49e39', 'folders': {'root': {'uuid': '486c9c58-7715-4bfc-a125-ec6ceaef5951', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '21ae4eed-489f-43a7-af0a-f70000714c79': {'uuid': '21ae4eed-489f-43a7-af0a-f70000714c79', 'hostname': 'web_server', 'operating_state': 2, 'NICs': {'060e96d7-399c-49d1-afcb-d448b3eec2d5': {'uuid': '060e96d7-399c-49d1-afcb-d448b3eec2d5', 'ip_adress': '192.168.1.12', 'subnet_mask': '255.255.255.0', 'mac_address': '7f:e0:b3:05:2c:d4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e8b1c5c5-f5ce-4d2f-840a-8b56e3fbd06b', 'folders': {'root': {'uuid': '86857c08-f5ed-4c60-98c9-a5c7f0985549', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'c4b7d388-c13e-4f4e-be61-5cf05989da71': {'uuid': 'c4b7d388-c13e-4f4e-be61-5cf05989da71', 'hostname': 'database_server', 'operating_state': 2, 'NICs': {'bed479cf-480b-463b-9fff-66e4dd3d23cc': {'uuid': 'bed479cf-480b-463b-9fff-66e4dd3d23cc', 'ip_adress': '192.168.1.14', 'subnet_mask': '255.255.255.0', 'mac_address': 'b4:d7:b6:04:7b:07', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '9c6f86b8-c14f-4593-8397-e3fe2e88ef66', 'folders': {'root': {'uuid': '1fd374f4-398a-4b98-9295-c2969cb3025c', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '12dc9e72-fb7c-4a8f-ab79-d7574863fa16': {'uuid': '12dc9e72-fb7c-4a8f-ab79-d7574863fa16', 'hostname': 'backup_server', 'operating_state': 2, 'NICs': {'b299a15c-cf03-4c50-92ed-ac2932472927': {'uuid': 'b299a15c-cf03-4c50-92ed-ac2932472927', 'ip_adress': '192.168.1.16', 'subnet_mask': '255.255.255.0', 'mac_address': '7b:0e:7c:b3:d1:9c', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '455a7e64-bee9-41e7-80c5-a38a26d3e7db', 'folders': {'root': {'uuid': 'd5f22eda-5f15-4df6-b9c5-5f107a0d7ef0', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '37ff728f-64b5-4203-8177-cf27f51dc7c9': {'uuid': '37ff728f-64b5-4203-8177-cf27f51dc7c9', 'hostname': 'security_suite', 'operating_state': 2, 'NICs': {'9011d432-dd95-4b0b-b793-2f9d173009da': {'uuid': '9011d432-dd95-4b0b-b793-2f9d173009da', 'ip_adress': '192.168.1.110', 'subnet_mask': '255.255.255.0', 'mac_address': '58:17:29:b0:d2:e4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '9bb869b9-7253-4890-92bc-b3f7e5119b6d', 'folders': {'root': {'uuid': '5415d598-e8ad-445d-b573-f4b56e6268e3', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a': {'uuid': 'fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a', 'hostname': 'client_1', 'operating_state': 2, 'NICs': {'80e94189-3bd7-45f5-ac04-50974e6db2e1': {'uuid': '80e94189-3bd7-45f5-ac04-50974e6db2e1', 'ip_adress': '192.168.10.21', 'subnet_mask': '255.255.255.0', 'mac_address': 'e1:ce:28:ef:74:76', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e4caeb03-81c1-4c0f-b619-f305ffd35c57', 'folders': {'root': {'uuid': 'b7f5c0c0-e41e-4d79-9ef8-d06e3e601929', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6': {'uuid': 'dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6', 'hostname': 'client_2', 'operating_state': 2, 'NICs': {'631bf440-da8f-41f0-947b-fdef122410ec': {'uuid': '631bf440-da8f-41f0-947b-fdef122410ec', 'ip_adress': '192.168.10.22', 'subnet_mask': '255.255.255.0', 'mac_address': '67:da:5c:11:c7:e4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'd4723cc6-32b9-4b21-8eab-3517ceb47130', 'folders': {'root': {'uuid': '2bbf4393-b5b3-4c2d-8b8e-6e6b883cdace', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}}, 'links': {}}, 'domain': {'uuid': '3216790f-d143-47e8-aa8d-8a2f819b34c7', 'accounts': {}}}\n" ] } ], diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 7870caab..2e7681a9 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -60,11 +60,11 @@ class ACLRule(SimComponent): """ state = super().describe_state() state["action"] = self.action.value - state["protocol"] = self.protocol.value - state["src_ip_address"] = self.src_ip_address - state["src_port"] = self.src_port.value - state["dst_ip_address"] = self.dst_ip_address - state["dst_port"] = self.dst_port.value + state["protocol"] = self.protocol.value if self.protocol else None + state["src_ip_address"] = self.src_ip_address if self.src_ip_address else None + state["src_port"] = self.src_port.value if self.src_port else None + state["dst_ip_address"] = self.dst_ip_address if self.dst_ip_address else None + state["dst_port"] = self.dst_port.value if self.dst_port else None return state From fdf66ba3de2c4b5085c55c9834ca9ff32e166c4b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 26 Sep 2023 10:52:14 +0100 Subject: [PATCH 06/35] Make yaml parser work with services --- example_config.yaml | 16 ++-- sandbox.ipynb | 219 ++++++++------------------------------------ 2 files changed, 48 insertions(+), 187 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index f0957718..dd5971a1 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -170,7 +170,7 @@ simulation: default_gateway: 192.168.1.1 services: - ref: domain_controller_dns_server - type: dns_server + type: DNSServer options: domain_mapping: - arcd.com: 192.168.1.12 # web server @@ -185,7 +185,7 @@ simulation: dns_server: 192.168.1.10 services: - ref: web_server_database_client - type: database_client + type: DatabaseClient options: db_server_ip: 192.168.1.14 @@ -198,7 +198,7 @@ simulation: dns_server: 192.168.1.10 services: - ref: database_service - type: database_service + type: DatabaseService - ref: backup_server @@ -210,7 +210,7 @@ simulation: dns_server: 192.168.1.10 services: - ref: backup_service - type: database_backup + type: DatabaseBackup - ref: security_suite type: server @@ -234,9 +234,9 @@ simulation: dns_server: 192.168.1.10 services: - ref: data_manipulation_bot - type: data_manipulation_bot + type: DataManipulationBot - ref: client_1_dns_client - type: dns_client + type: DNSClient - ref: client_2 type: computer @@ -247,9 +247,9 @@ simulation: dns_server: 192.168.1.10 services: - ref: client_2_web_browser - type: web_browser + type: WebBrowser - ref: client_2_dns_client - type: dns_client + type: DNSClient links: diff --git a/sandbox.ipynb b/sandbox.ipynb index 96d12bae..5d611ada 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -12,152 +12,7 @@ }, { "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "import yaml" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from primaite.simulator.network.networks import arcd_uc2_network\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "net = arcd_uc2_network()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "random_node = list(net.nodes.keys())[0]\n", - "f = net.nodes[random_node].file_system.create_file(file_name=\"testfile\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "f.describe_state()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def test_file_observation():\n", - " from primaite.simulator.sim_container import Simulation\n", - " from primaite.simulator.network.hardware.nodes.computer import Computer\n", - " from primaite.game.actor.observations import FileObservation\n", - "\n", - " sim = Simulation()\n", - " pc = Computer(hostname=\"beep\", ip_address=\"123.123.123.123\", subnet_mask=\"255.255.255.0\")\n", - " sim.network.add_node(pc)\n", - " f = pc.file_system.create_file(file_name=\"dog.png\")\n", - "\n", - " dog_file_obs = FileObservation(where=['network','nodes',pc.uuid,'file_system'])\n", - " print(sim.describe_state())\n", - "test_file_observation()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from primaite.simulator.sim_container import Simulation\n", - "from primaite.simulator.network.hardware.nodes.computer import Computer\n", - "from primaite.game.actor.observations import FileObservation, FolderObservation\n", - "\n", - "sim = Simulation()\n", - "pc = Computer(hostname=\"beep\", ip_address=\"123.123.123.123\", subnet_mask=\"255.255.255.0\")\n", - "sim.network.add_node(pc)\n", - "f = pc.file_system.create_file(file_name=\"dog.png\")\n", - "\n", - "state = sim.describe_state()\n", - "\n", - "dog_file_obs = FileObservation(where=['network','nodes',pc.uuid,'file_system', 'folders','root','files','dog.png'])\n", - "root_folder_obs = FolderObservation(where=['network','nodes',pc.uuid,'file_system', 'folders','root'],files=[dog_file_obs])\n", - "print(dog_file_obs(state))\n", - "print(root_folder_obs(state))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dog_file_obs.space" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "root_folder_obs.space" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "state" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import yaml" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "with open('example_config.yaml', 'r') as file:\n", - " conf = yaml.safe_load(file)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "conf['simulation']" - ] - }, - { - "cell_type": "code", - "execution_count": 44, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -172,39 +27,48 @@ "from primaite.simulator.system.services.database_service import DatabaseService\n", "from primaite.simulator.system.services.dns_client import DNSClient\n", "from primaite.simulator.system.services.dns_server import DNSServer\n", - "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot\n" + "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot\n", + "\n", + "\n", + "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", + "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", + "from primaite.simulator.network.transmission.transport_layer import Port\n", + "\n" ] }, { "cell_type": "code", - "execution_count": 42, + "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-25 19:10:46,253: Added node 3d356af7-15f8-41d4-bb2a-423d5a2e5978 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,254::WARNING::primaite.simulator.network.container::181::Can't add node 3d356af7-15f8-41d4-bb2a-423d5a2e5978. It is already in the network.\n", - "2023-09-25 19:10:46,258: Added node 8a94447a-ccb3-47b2-b7ba-488e631f8246 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,262: Added node 4e3de72a-60ad-416f-abb7-da352a59d13b to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,264: Added node 21ae4eed-489f-43a7-af0a-f70000714c79 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,267: Added node c4b7d388-c13e-4f4e-be61-5cf05989da71 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,270: Added node 12dc9e72-fb7c-4a8f-ab79-d7574863fa16 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,273: Added node 37ff728f-64b5-4203-8177-cf27f51dc7c9 to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,277: Added node fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a to Network 44950e57-81b0-4964-9b12-223592c785aa\n", - "2023-09-25 19:10:46,281: Added node dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6 to Network 44950e57-81b0-4964-9b12-223592c785aa\n" + "2023-09-26 10:51:10,388: Added node 48e6cb0b-f351-47f6-b837-df9443f9db26 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,390::WARNING::primaite.simulator.network.container::181::Can't add node 48e6cb0b-f351-47f6-b837-df9443f9db26. It is already in the network.\n", + "2023-09-26 10:51:10,394: Added node 6a969d4d-e0af-402e-b576-2a787505f7c7 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,397: Added node c58e6f17-dbf1-4c6a-9dbf-d60883c6d948 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,401: Added node 7f2a418d-2d0b-4f02-beb3-5703fc5035c8 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,408: Added node 967417fa-2300-4ee1-8ba0-7a4d055d5d30 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,413: Added node 80c1c99b-4c7a-41fb-86f0-b93c35c3b497 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,418: Added node 9a11dd40-9243-4510-9b43-9f247f669ad2 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,424: Added node 81fff4a6-35c8-4933-bb6c-fd8fd49315fe to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", + "2023-09-26 10:51:10,429: Added node 7cc11532-3f65-4c65-a4df-af2c6318a976 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "service type not found DatabaseBackup\n", + "service type not found WebBrowser\n" ] } ], "source": [ "# import yaml\n", "\n", - "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", - "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", - "from primaite.simulator.network.transmission.transport_layer import Port\n", - "from primaite.simulator.system.services.dns_server import DNSServer\n", - "\n", "\n", "class PrimaiteSession:\n", "\n", @@ -270,15 +134,20 @@ " service_ref = service_cfg['ref']\n", " service_type = service_cfg['type']\n", " service_types_mapping = {\n", - " 'dns_server' : DNSServer,\n", - " 'database_client': DatabaseClient,\n", - " 'database_service': DatabaseService,\n", + " 'DNSClient': DNSClient, # key is equal to the 'name' attr of the service class itself.\n", + " 'DNSServer' : DNSServer,\n", + " 'DatabaseClient': DatabaseClient,\n", + " 'DatabaseService': DatabaseService,\n", " # 'database_backup': ,\n", - " 'data_manipulation_bot': DataManipulationBot,\n", - " 'dns_client': DNSClient,\n", + " 'DataManipulationBot': DataManipulationBot,\n", " # 'web_browser'\n", " }\n", - " if service_type\n", + " if service_type in service_types_mapping:\n", + " new_node.software_manager.install(service_types_mapping[service_type])\n", + " service_obj = new_node.software_manager.software[service_type]\n", + " ref_to_uuid[service_ref] = service_obj.uuid\n", + " else:\n", + " print(f\"service type not found {service_type}\")\n", "\n", "\n", " net.add_node(new_node)\n", @@ -297,17 +166,9 @@ }, { "cell_type": "code", - "execution_count": 43, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'uuid': '532b534a-78b7-4630-92f9-626d775ee4e2', 'network': {'uuid': '44950e57-81b0-4964-9b12-223592c785aa', 'nodes': {'3d356af7-15f8-41d4-bb2a-423d5a2e5978': {'uuid': '3d356af7-15f8-41d4-bb2a-423d5a2e5978', 'hostname': 'router_1', 'operating_state': 2, 'NICs': {'c52a1516-3488-4011-95ae-44e77935720d': {'uuid': 'c52a1516-3488-4011-95ae-44e77935720d', 'ip_adress': '192.168.1.1', 'subnet_mask': '255.255.255.0', 'mac_address': '12:75:33:08:ae:04', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '075a3387-b3c5-4e30-b7d5-f39b82b1c9d8': {'uuid': '075a3387-b3c5-4e30-b7d5-f39b82b1c9d8', 'ip_adress': '192.168.1.1', 'subnet_mask': '255.255.255.0', 'mac_address': 'ce:67:79:4a:48:90', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '3fe76cb7-84fe-4592-b527-be69c06c8064': {'uuid': '3fe76cb7-84fe-4592-b527-be69c06c8064', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '83:13:43:3a:e6:de', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, 'b4ad54f8-8508-4d34-9117-8f194a1076b1': {'uuid': 'b4ad54f8-8508-4d34-9117-8f194a1076b1', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '4f:1a:3f:76:c7:2c', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}, '21efd6b1-8b49-4b56-aafe-a3cbdf61e0fb': {'uuid': '21efd6b1-8b49-4b56-aafe-a3cbdf61e0fb', 'ip_adress': '127.0.0.1', 'subnet_mask': '255.0.0.0', 'mac_address': '76:4d:65:e4:0d:11', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'df3b627a-76c3-4e60-83ef-81442fbf4643', 'folders': {'root': {'uuid': 'b21a0c20-e852-4d05-a0fe-9158f76f3a98', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}, 'num_ports': (5,), 'acl': ({'uuid': '2fc154d6-9ccf-4c60-b6d3-ac3fcce68fc8', 'implicit_action': 0, 'implicit_rule': {'uuid': '6a5553d0-37a6-4ac5-b5da-38f7926ffe7e', 'action': 0, 'protocol': None, 'src_ip_address': None, 'src_port': None, 'dst_ip_address': None, 'dst_port': None}, 'max_acl_rules': 25, 'acl': {0: {'uuid': '8204a738-6010-44a8-8df3-317d4a03f600', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 5432, 'dst_ip_address': None, 'dst_port': 5432}, 1: {'uuid': 'af7b39a5-e0ed-4047-b684-0cfbdc14f70b', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 53, 'dst_ip_address': None, 'dst_port': 53}, 2: None, 3: None, 4: None, 5: None, 6: None, 7: None, 8: None, 9: None, 10: None, 11: None, 12: None, 13: None, 14: None, 15: None, 16: None, 17: None, 18: None, 19: None, 20: None, 21: None, 22: {'uuid': '4c6320f5-d9f4-4958-9572-d2ebc06ee120', 'action': 1, 'protocol': None, 'src_ip_address': None, 'src_port': 219, 'dst_ip_address': None, 'dst_port': 219}, 23: {'uuid': 'f39796b9-7511-4f23-b8a6-cbf0e87ec51a', 'action': 1, 'protocol': 'icmp', 'src_ip_address': None, 'src_port': None, 'dst_ip_address': None, 'dst_port': None}}},)}, '8a94447a-ccb3-47b2-b7ba-488e631f8246': {'uuid': '8a94447a-ccb3-47b2-b7ba-488e631f8246', 'num_ports': 8, 'ports': {1: {'uuid': '32eb0312-124d-42df-93bc-9cc25eb4b81f', 'mac_address': '1a:4e:83:7a:d7:4b', 'speed': 100, 'mtu': 1500, 'enabled': False}, 2: {'uuid': 'afe7a1b8-bfa6-42a2-b4aa-28187c67f7d3', 'mac_address': '9e:09:af:37:1d:e3', 'speed': 100, 'mtu': 1500, 'enabled': False}, 3: {'uuid': '98f46066-7e05-45cd-920e-54bbb3513fec', 'mac_address': '00:30:02:18:11:52', 'speed': 100, 'mtu': 1500, 'enabled': False}, 4: {'uuid': '65fb27f3-22ca-4bd3-898e-a2fdf49ec3b0', 'mac_address': 'f2:3f:0b:09:50:e1', 'speed': 100, 'mtu': 1500, 'enabled': False}, 5: {'uuid': 'f69c7045-40f7-490f-9c7c-280d0eb4fcd0', 'mac_address': '38:56:f5:dd:0a:a4', 'speed': 100, 'mtu': 1500, 'enabled': False}, 6: {'uuid': '6dbbac83-cd01-45ec-9ad6-352e1afee231', 'mac_address': 'aa:bb:c3:92:05:88', 'speed': 100, 'mtu': 1500, 'enabled': False}, 7: {'uuid': 'a39591ea-d50d-4649-ba90-f5fddd55b61d', 'mac_address': '60:67:46:45:21:cc', 'speed': 100, 'mtu': 1500, 'enabled': False}, 8: {'uuid': '488cdddb-9e5e-4fb1-8e03-cfcfd47b0cb3', 'mac_address': '0f:a0:4f:52:42:5c', 'speed': 100, 'mtu': 1500, 'enabled': False}}, 'mac_address_table': {}}, '4e3de72a-60ad-416f-abb7-da352a59d13b': {'uuid': '4e3de72a-60ad-416f-abb7-da352a59d13b', 'hostname': 'domain_controller', 'operating_state': 2, 'NICs': {'34962c55-cea2-479b-ac1f-61a683c82eb1': {'uuid': '34962c55-cea2-479b-ac1f-61a683c82eb1', 'ip_adress': '192.168.1.10', 'subnet_mask': '255.255.255.0', 'mac_address': '41:a0:d4:22:b9:81', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e3f270c9-2941-4592-892c-5662bbc49e39', 'folders': {'root': {'uuid': '486c9c58-7715-4bfc-a125-ec6ceaef5951', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '21ae4eed-489f-43a7-af0a-f70000714c79': {'uuid': '21ae4eed-489f-43a7-af0a-f70000714c79', 'hostname': 'web_server', 'operating_state': 2, 'NICs': {'060e96d7-399c-49d1-afcb-d448b3eec2d5': {'uuid': '060e96d7-399c-49d1-afcb-d448b3eec2d5', 'ip_adress': '192.168.1.12', 'subnet_mask': '255.255.255.0', 'mac_address': '7f:e0:b3:05:2c:d4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e8b1c5c5-f5ce-4d2f-840a-8b56e3fbd06b', 'folders': {'root': {'uuid': '86857c08-f5ed-4c60-98c9-a5c7f0985549', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'c4b7d388-c13e-4f4e-be61-5cf05989da71': {'uuid': 'c4b7d388-c13e-4f4e-be61-5cf05989da71', 'hostname': 'database_server', 'operating_state': 2, 'NICs': {'bed479cf-480b-463b-9fff-66e4dd3d23cc': {'uuid': 'bed479cf-480b-463b-9fff-66e4dd3d23cc', 'ip_adress': '192.168.1.14', 'subnet_mask': '255.255.255.0', 'mac_address': 'b4:d7:b6:04:7b:07', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '9c6f86b8-c14f-4593-8397-e3fe2e88ef66', 'folders': {'root': {'uuid': '1fd374f4-398a-4b98-9295-c2969cb3025c', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '12dc9e72-fb7c-4a8f-ab79-d7574863fa16': {'uuid': '12dc9e72-fb7c-4a8f-ab79-d7574863fa16', 'hostname': 'backup_server', 'operating_state': 2, 'NICs': {'b299a15c-cf03-4c50-92ed-ac2932472927': {'uuid': 'b299a15c-cf03-4c50-92ed-ac2932472927', 'ip_adress': '192.168.1.16', 'subnet_mask': '255.255.255.0', 'mac_address': '7b:0e:7c:b3:d1:9c', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '455a7e64-bee9-41e7-80c5-a38a26d3e7db', 'folders': {'root': {'uuid': 'd5f22eda-5f15-4df6-b9c5-5f107a0d7ef0', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, '37ff728f-64b5-4203-8177-cf27f51dc7c9': {'uuid': '37ff728f-64b5-4203-8177-cf27f51dc7c9', 'hostname': 'security_suite', 'operating_state': 2, 'NICs': {'9011d432-dd95-4b0b-b793-2f9d173009da': {'uuid': '9011d432-dd95-4b0b-b793-2f9d173009da', 'ip_adress': '192.168.1.110', 'subnet_mask': '255.255.255.0', 'mac_address': '58:17:29:b0:d2:e4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': '9bb869b9-7253-4890-92bc-b3f7e5119b6d', 'folders': {'root': {'uuid': '5415d598-e8ad-445d-b573-f4b56e6268e3', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a': {'uuid': 'fe9ef0ad-e0d6-48b9-a884-f7c0b95de32a', 'hostname': 'client_1', 'operating_state': 2, 'NICs': {'80e94189-3bd7-45f5-ac04-50974e6db2e1': {'uuid': '80e94189-3bd7-45f5-ac04-50974e6db2e1', 'ip_adress': '192.168.10.21', 'subnet_mask': '255.255.255.0', 'mac_address': 'e1:ce:28:ef:74:76', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'e4caeb03-81c1-4c0f-b619-f305ffd35c57', 'folders': {'root': {'uuid': 'b7f5c0c0-e41e-4d79-9ef8-d06e3e601929', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}, 'dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6': {'uuid': 'dfd186b9-7dc7-4d6f-a736-0cce3d22bfb6', 'hostname': 'client_2', 'operating_state': 2, 'NICs': {'631bf440-da8f-41f0-947b-fdef122410ec': {'uuid': '631bf440-da8f-41f0-947b-fdef122410ec', 'ip_adress': '192.168.10.22', 'subnet_mask': '255.255.255.0', 'mac_address': '67:da:5c:11:c7:e4', 'speed': 100, 'mtu': 1500, 'wake_on_lan': False, 'enabled': False}}, 'file_system': {'uuid': 'd4723cc6-32b9-4b21-8eab-3517ceb47130', 'folders': {'root': {'uuid': '2bbf4393-b5b3-4c2d-8b8e-6e6b883cdace', 'name': 'root', 'health_status': 1, 'files': {}, 'is_quarantined': False}}}, 'applications': {}, 'services': {}, 'process': {}}}, 'links': {}}, 'domain': {'uuid': '3216790f-d143-47e8-aa8d-8a2f819b34c7', 'accounts': {}}}\n" - ] - } - ], + "outputs": [], "source": [ "print(s.simulation.describe_state())" ] From 92e0110e73de13e022adb8a55c434d9207f28c28 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 26 Sep 2023 11:48:22 +0100 Subject: [PATCH 07/35] yaml parse and connect links --- example_config.yaml | 80 ++++++++++----------- sandbox.ipynb | 74 ++++++++++++++----- src/primaite/simulator/network/container.py | 3 +- 3 files changed, 99 insertions(+), 58 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index dd5971a1..e3871b4a 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -22,12 +22,12 @@ game_config: actions: - type: DONOTHING nodes: - - ref: client_2 + - node_ref: client_2 actions: - type: LOGON - type: LOGOFF applications: - - ref: client_2_web_browser + - application_ref: client_2_web_browser actions: - type: EXECUTE execution_definition: @@ -44,12 +44,12 @@ game_config: observation_space: network: nodes: - - ref: client_1 + - node_ref: client_1 observations: - logon_status - operating_status services: - - ref: data_manipulation_bot + - service_ref: data_manipulation_bot observations: - operating_status - health_status @@ -59,13 +59,13 @@ game_config: - type: DO_NOTHING network: nodes: - - ref: client_1 + - node_ref: client_1 actions: - type: SCAN - type: LOGON - type: LOGOFF services: - - ref: data_manipulation_bot + - service_ref: data_manipulation_bot actions: - type: COMPROMISE execution_definition: @@ -89,18 +89,18 @@ game_config: observation_space: network: nodes: - - ref: router_1 #TODO: more sub-options here - - ref: switch_1 - - ref: switch_2 - - ref: domain_controller - - ref: web_server - - ref: database_server - - ref: backup_server - - ref: security_suite - - ref: client_1 - - ref: client_2 + - node_ref: router_1 #TODO: more sub-options here + - node_ref: switch_1 + - node_ref: switch_2 + - node_ref: domain_controller + - node_ref: web_server + - node_ref: database_server + - node_ref: backup_server + - node_ref: security_suite + - node_ref: client_1 + - node_ref: client_2 links: - - ref: ... # + - link_ref: ... # acl: ... # ics: ... # @@ -110,7 +110,7 @@ game_config: - type: DO_NOTHING network: nodes: - - ref: router_1 + - node_ref: router_1 reward_function: # ... agent_settings: @@ -153,7 +153,7 @@ simulation: protocol: ICMP - ref: switch_1 - type: swtich + type: switch hostname: switch_1 num_ports: 8 @@ -173,7 +173,7 @@ simulation: type: DNSServer options: domain_mapping: - - arcd.com: 192.168.1.12 # web server + arcd.com: 192.168.1.12 # web server - ref: web_server @@ -254,52 +254,52 @@ simulation: links: - ref: router_1___switch_1 - endpoint_a: router_1 + endpoint_a_ref: router_1 endpoint_a_port: 1 - endpoint_b: switch_1 + endpoint_b_ref: switch_1 endpoint_b_port: 8 - ref: router_1___switch_2 - endpoint_a: router_1 + endpoint_a_ref: router_1 endpoint_a_port: 2 - endpoint_b: switch_2 + endpoint_b_ref: switch_2 endpoint_b_port: 8 - ref: switch_1___domain_controller - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 1 - endpoint_b: domain_controller + endpoint_b_ref: domain_controller endpoint_b_port: 1 - ref: switch_1___web_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 2 - endpoint_b: web_server + endpoint_b_ref: web_server endpoint_b_port: 1 - ref: switch_1___database_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 3 - endpoint_b: database_server + endpoint_b_ref: database_server endpoint_b_port: 1 - ref: switch_1___backup_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 4 - endpoint_b: backup_server + endpoint_b_ref: backup_server endpoint_b_port: 1 - ref: switch_1___security_suite - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 7 - endpoint_b: security_suite + endpoint_b_ref: security_suite endpoint_b_port: 1 - ref: switch_2___client_1 - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 1 - endpoint_b: client_1 + endpoint_b_ref: client_1 endpoint_b_port: 1 - ref: switch_2___client_2 - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 2 - endpoint_b: client_2 + endpoint_b_ref: client_2 endpoint_b_port: 1 - ref: switch_2___security_suite - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 7 - endpoint_b: security_suite + endpoint_b_ref: security_suite endpoint_b_port: 2 diff --git a/sandbox.ipynb b/sandbox.ipynb index 5d611ada..aa39c3e9 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -33,28 +33,29 @@ "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", "from primaite.simulator.network.transmission.transport_layer import Port\n", - "\n" + "\n", + "from ipaddress import IPv4Address\n" ] }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-26 10:51:10,388: Added node 48e6cb0b-f351-47f6-b837-df9443f9db26 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,390::WARNING::primaite.simulator.network.container::181::Can't add node 48e6cb0b-f351-47f6-b837-df9443f9db26. It is already in the network.\n", - "2023-09-26 10:51:10,394: Added node 6a969d4d-e0af-402e-b576-2a787505f7c7 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,397: Added node c58e6f17-dbf1-4c6a-9dbf-d60883c6d948 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,401: Added node 7f2a418d-2d0b-4f02-beb3-5703fc5035c8 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,408: Added node 967417fa-2300-4ee1-8ba0-7a4d055d5d30 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,413: Added node 80c1c99b-4c7a-41fb-86f0-b93c35c3b497 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,418: Added node 9a11dd40-9243-4510-9b43-9f247f669ad2 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,424: Added node 81fff4a6-35c8-4933-bb6c-fd8fd49315fe to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,429: Added node 7cc11532-3f65-4c65-a4df-af2c6318a976 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n" + "2023-09-26 11:47:11,032: Added node bc149bf5-ccc4-4dcd-b419-629ec44b2c9a to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,035: Added node 9cacbaee-33cc-4423-a6c8-fe3dd75b1f87 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,042: Added node d4444d66-7cc3-4cd4-acbd-202cb9fe37ff to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,045: Added node af170371-e99b-42b7-9525-65ca64522539 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,049: Added node d6218f34-a104-469d-a08b-97329ad84c19 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,052: Added node 831a3803-ae65-4cee-a17e-9c1220035bc9 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,055: Added node 1b935654-065d-4cb9-82d9-d67fe3d3304e to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,059: Added node dd181916-076b-4d8a-ab97-a32052624b09 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,064: Added node 3137ab20-1a3c-49f2-8ee5-c862216b2435 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,067: Added node 6ff8b634-7750-4c6d-8109-abf52514dae5 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n" ] }, { @@ -70,6 +71,11 @@ "# import yaml\n", "\n", "\n", + "from typing import Dict\n", + "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", + "from primaite.simulator.system.services.service import Service\n", + "\n", + "\n", "class PrimaiteSession:\n", "\n", " def __init__(self):\n", @@ -78,7 +84,11 @@ "\n", " @classmethod\n", " def from_config(cls, cfg_path):\n", - " ref_to_uuid = {}\n", + " ref_map_nodes: Dict[str,Node] = {}\n", + " ref_map_services: Dict[str, Service] = {}\n", + " ref_map_links: Dict[str, Link] = {}\n", + " # ref_map_agents: Dict[str, AgentInterface] = {}\n", + "\n", "\n", " game = cls()\n", " with open(cfg_path, 'r') as file:\n", @@ -129,6 +139,8 @@ " dst_ip_address = r_cfg.get('ip_address'),\n", " position = r_num\n", " )\n", + " else:\n", + " print('invalid node type')\n", " if 'services' in node_cfg:\n", " for service_cfg in node_cfg['services']:\n", " service_ref = service_cfg['ref']\n", @@ -144,15 +156,43 @@ " }\n", " if service_type in service_types_mapping:\n", " new_node.software_manager.install(service_types_mapping[service_type])\n", - " service_obj = new_node.software_manager.software[service_type]\n", - " ref_to_uuid[service_ref] = service_obj.uuid\n", + " new_service = new_node.software_manager.software[service_type]\n", + " ref_map_services[service_ref] = new_service\n", " else:\n", " print(f\"service type not found {service_type}\")\n", - "\n", + " # service-dependent options\n", + " if service_type == 'DatabaseClient':\n", + " if 'options' in service_cfg:\n", + " opt = service_cfg['options']\n", + " if 'db_server_ip' in opt:\n", + " new_service.configure(server_ip_address=IPv4Address(opt['db_server_ip']))\n", + " if service_type == 'DNSServer':\n", + " if 'options' in service_cfg:\n", + " opt = service_cfg['options']\n", + " if 'domain_mapping' in opt:\n", + " for domain, ip in opt['domain_mapping'].items():\n", + " new_service.dns_register(domain, ip)\n", + " if 'nics' in node_cfg:\n", + " for nic_num, nic_cfg in node_cfg['nics'].items():\n", + " new_node.connect_nic(NIC(ip_address=nic_cfg['ip_address'], subnet_mask=nic_cfg['subnet_mask']))\n", "\n", " net.add_node(new_node)\n", - " ref_to_uuid[node_ref] = new_node.uuid\n", + " ref_map_nodes[node_ref] = new_node.uuid\n", "\n", + " #2. create links between nodes\n", + " for link_cfg in links_cfg:\n", + " node_a = net.nodes[ref_map_nodes[link_cfg['endpoint_a_ref']]]\n", + " node_b = net.nodes[ref_map_nodes[link_cfg['endpoint_b_ref']]]\n", + " if isinstance(node_a, Switch):\n", + " endpoint_a = node_a.switch_ports[link_cfg['endpoint_a_port']]\n", + " else:\n", + " endpoint_a = node_a.ethernet_port[link_cfg['endpoint_a_port']]\n", + " if isinstance(node_b, Switch):\n", + " endpoint_b = node_b.switch_ports[link_cfg['endpoint_b_port']]\n", + " else:\n", + " endpoint_b = node_b.ethernet_port[link_cfg['endpoint_b_port']]\n", + " new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)\n", + " ref_map_links[link_cfg['ref']] = new_link.uuid\n", "\n", " #2. start/setup simulation objects\n", " #3. create agents\n", diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 7ab9b093..66686797 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -221,7 +221,7 @@ class Network(SimComponent): _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}") self._node_action_manager.remove_action(name=node.uuid) - def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None: + def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> Optional[Link]: """ Connect two endpoints on the network by creating a link between their NICs/SwitchPorts. @@ -248,6 +248,7 @@ class Network(SimComponent): self._nx_graph.add_edge(endpoint_a.parent.hostname, endpoint_b.parent.hostname) link.parent = self _LOGGER.debug(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}") + return link def remove_link(self, link: Link) -> None: """Disconnect a link from the network. From f1346ae278e8659ce0ef4cc15daced2ad88d2ff2 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 26 Sep 2023 12:54:56 +0100 Subject: [PATCH 08/35] put in agent parsing skeleton --- example_config.yaml | 6 +-- sandbox.ipynb | 48 +++++++++++++------ src/primaite/game/actor/interface.py | 35 -------------- src/primaite/game/agent/GATE_agents.py | 5 ++ .../game/{actor => agent}/__init__.py | 0 src/primaite/game/{actor => agent}/actions.py | 0 src/primaite/game/agent/interface.py | 35 ++++++++++++++ .../game/{actor => agent}/observations.py | 0 src/primaite/game/{actor => agent}/rewards.py | 0 src/primaite/game/agent/scripted_agents.py | 9 ++++ .../game_layer/test_observations.py | 2 +- 11 files changed, 86 insertions(+), 54 deletions(-) delete mode 100644 src/primaite/game/actor/interface.py create mode 100644 src/primaite/game/agent/GATE_agents.py rename src/primaite/game/{actor => agent}/__init__.py (100%) rename src/primaite/game/{actor => agent}/actions.py (100%) create mode 100644 src/primaite/game/agent/interface.py rename src/primaite/game/{actor => agent}/observations.py (100%) rename src/primaite/game/{actor => agent}/rewards.py (100%) create mode 100644 src/primaite/game/agent/scripted_agents.py diff --git a/example_config.yaml b/example_config.yaml index e3871b4a..79cfccac 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -16,7 +16,7 @@ game_config: agents: - ref: client_1_green_user team: GREEN - team: SCRIPTED_GREEN_ + type: GreenWebBrowsingAgent observation_space: null action_space: actions: @@ -40,7 +40,7 @@ game_config: - ref: client_1_data_manipulation_red_bot team: RED - type: SCRIPTED_RED_ + type: RedDatabaseCorruptingAgent observation_space: network: nodes: @@ -220,7 +220,7 @@ simulation: default_gateway: 192.168.1.1 dns_server: 192.168.1.10 nics: - 2: + 2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot ip_address: 192.168.10.110 subnet_mask: 255.255.255.0 diff --git a/sandbox.ipynb b/sandbox.ipynb index aa39c3e9..05efcfa2 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -39,23 +39,23 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 18, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-26 11:47:11,032: Added node bc149bf5-ccc4-4dcd-b419-629ec44b2c9a to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,035: Added node 9cacbaee-33cc-4423-a6c8-fe3dd75b1f87 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,042: Added node d4444d66-7cc3-4cd4-acbd-202cb9fe37ff to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,045: Added node af170371-e99b-42b7-9525-65ca64522539 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,049: Added node d6218f34-a104-469d-a08b-97329ad84c19 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,052: Added node 831a3803-ae65-4cee-a17e-9c1220035bc9 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,055: Added node 1b935654-065d-4cb9-82d9-d67fe3d3304e to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,059: Added node dd181916-076b-4d8a-ab97-a32052624b09 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,064: Added node 3137ab20-1a3c-49f2-8ee5-c862216b2435 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,067: Added node 6ff8b634-7750-4c6d-8109-abf52514dae5 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n" + "2023-09-26 12:19:50,895: Added node 0fb262e1-a714-420a-aec7-be37f0deeb75 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,898: Added node 310ca8d7-01e0-401e-b604-705c290e5376 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,900: Added node b3b08f1f-7805-47b2-bdb6-3d83098cd740 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,903: Added node adb37f3e-2307-4123-bff3-01f125883be8 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,906: Added node 1a490716-2ccd-452d-b87e-324d29120b59 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,911: Added node 033460d8-0249-4bdd-aaf0-751b24cc0a1e to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,914: Added node 1e7e4e49-78bf-4031-8372-ee71902720f3 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,916: Added node c9f24a13-e5c8-437b-9234-b0c3f8120e2c to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,920: Added node c881f3ee-2176-493b-a6c2-cad829bf0b6d to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,922: Added node a3ea75d8-bc2c-4713-92a4-2588b4f43ed6 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n" ] }, { @@ -72,9 +72,12 @@ "\n", "\n", "from typing import Dict\n", + "from primaite.game.agent.interface import AbstractAgent\n", "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", "from primaite.simulator.system.services.service import Service\n", "\n", + "from primaite.game.agent.scripted_agents import GreenWebBrowsingAgent, RedDatabaseCorruptingAgent\n", + "from primaite.game.agent.GATE_agents import GATERLAgent\n", "\n", "class PrimaiteSession:\n", "\n", @@ -90,7 +93,7 @@ " # ref_map_agents: Dict[str, AgentInterface] = {}\n", "\n", "\n", - " game = cls()\n", + " session = cls()\n", " with open(cfg_path, 'r') as file:\n", " conf = yaml.safe_load(file)\n", " \n", @@ -177,6 +180,7 @@ " new_node.connect_nic(NIC(ip_address=nic_cfg['ip_address'], subnet_mask=nic_cfg['subnet_mask']))\n", "\n", " net.add_node(new_node)\n", + " new_node.power_on()\n", " ref_map_nodes[node_ref] = new_node.uuid\n", "\n", " #2. create links between nodes\n", @@ -194,11 +198,25 @@ " new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)\n", " ref_map_links[link_cfg['ref']] = new_link.uuid\n", "\n", - " #2. start/setup simulation objects\n", + " session.simulation = sim\n", " #3. create agents\n", + " game_cfg = conf['game_config']\n", + " ports_cfg = game_cfg['ports']\n", + " protocols_cfg = game_cfg['protocols']\n", + " agents_cfg = game_cfg['agents']\n", + "\n", + " for agent_cfg in agents_cfg:\n", + " agent_ref = agent_cfg['ref']\n", + " agent_type = agent_cfg['type']\n", + " action_space_cfg = agent_cfg['action_space']\n", + " observation_space_cfg = agent_cfg['observation_space']\n", + " reward_function_cfg = agent_cfg['reward_function']\n", + " if agent_type == 'GreenWebBrowsingAgent':\n", + " new_agent = GreenWebBrowsingAgent()\n", + "\n", + "\n", " #4. set up agents' actions and observation spaces.\n", - " game.simulation = sim\n", - " return game\n", + " return session\n", "\n", "s = PrimaiteSession.from_config('example_config.yaml')\n", "# print(s.simulation.describe_state())" diff --git a/src/primaite/game/actor/interface.py b/src/primaite/game/actor/interface.py deleted file mode 100644 index d1245e71..00000000 --- a/src/primaite/game/actor/interface.py +++ /dev/null @@ -1,35 +0,0 @@ -# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' -# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word -# If you disagree, make a comment in the PR review and we can discuss -from abc import ABC, abstractmethod -from typing import Any, Dict, List - -from pydantic import BaseModel - -from primaite.game.actor.actions import ActionSpace -from primaite.game.actor.observations import ObservationSpace -from primaite.game.actor.rewards import RewardFunction - - -class AbstractActor(ABC): - """Base class for scripted and RL agents.""" - - def __init__(self) -> None: - self.action_space = ActionSpace - self.observation_space = ObservationSpace - self.reward_function = RewardFunction - - -class AbstractScriptedActor(AbstractActor): - """Base class for actors which generate their own behaviour.""" - - ... - - -class AbstractPuppetActor(AbstractActor): - """Base class for actors controlled via external messages, such as RL policies.""" - - ... - - -# class AbstractRLActor(AbstractPuppetActor): ?? diff --git a/src/primaite/game/agent/GATE_agents.py b/src/primaite/game/agent/GATE_agents.py new file mode 100644 index 00000000..5bdfebe4 --- /dev/null +++ b/src/primaite/game/agent/GATE_agents.py @@ -0,0 +1,5 @@ +from primaite.game.agent.interface import AbstractGATEAgent + + +class GATERLAgent(AbstractGATEAgent): + ... diff --git a/src/primaite/game/actor/__init__.py b/src/primaite/game/agent/__init__.py similarity index 100% rename from src/primaite/game/actor/__init__.py rename to src/primaite/game/agent/__init__.py diff --git a/src/primaite/game/actor/actions.py b/src/primaite/game/agent/actions.py similarity index 100% rename from src/primaite/game/actor/actions.py rename to src/primaite/game/agent/actions.py diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py new file mode 100644 index 00000000..b1ade94b --- /dev/null +++ b/src/primaite/game/agent/interface.py @@ -0,0 +1,35 @@ +# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' +# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word +# If you disagree, make a comment in the PR review and we can discuss +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +from primaite.game.agent.actions import ActionSpace +from primaite.game.agent.observations import ObservationSpace +from primaite.game.agent.rewards import RewardFunction + + +class AbstractAgent(ABC): + """Base class for scripted and RL agents.""" + + def __init__( + self, + action_space: Optional[ActionSpace], + observation_space: Optional[ObservationSpace], + reward_function: Optional[RewardFunction], + ) -> None: + self.action_space: Optional[ActionSpace] = action_space + self.observation_space: Optional[ObservationSpace] = observation_space + self.reward_function: Optional[RewardFunction] = reward_function + + +class AbstractScriptedAgent(AbstractAgent): + """Base class for actors which generate their own behaviour.""" + + ... + + +class AbstractGATEAgent(AbstractAgent): + """Base class for actors controlled via external messages, such as RL policies.""" + + ... diff --git a/src/primaite/game/actor/observations.py b/src/primaite/game/agent/observations.py similarity index 100% rename from src/primaite/game/actor/observations.py rename to src/primaite/game/agent/observations.py diff --git a/src/primaite/game/actor/rewards.py b/src/primaite/game/agent/rewards.py similarity index 100% rename from src/primaite/game/actor/rewards.py rename to src/primaite/game/agent/rewards.py diff --git a/src/primaite/game/agent/scripted_agents.py b/src/primaite/game/agent/scripted_agents.py new file mode 100644 index 00000000..d3becd57 --- /dev/null +++ b/src/primaite/game/agent/scripted_agents.py @@ -0,0 +1,9 @@ +from primaite.game.agent.interface import AbstractScriptedAgent + + +class GreenWebBrowsingAgent(AbstractScriptedAgent): + ... + + +class RedDatabaseCorruptingAgent(AbstractScriptedAgent): + ... diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index 4f778f78..7f20a938 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -1,6 +1,6 @@ from gym import spaces -from primaite.game.actor.observations import FileObservation +from primaite.game.agent.observations import FileObservation from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.sim_container import Simulation From 2b617e01a39263a64701f33f99e8773b28c1eac6 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 2 Oct 2023 17:21:43 +0100 Subject: [PATCH 09/35] Finalise actions interface --- example_config.yaml | 330 ++++++++++++++-- sandbox.ipynb | 489 ++++++++++++++++++++++-- src/primaite/game/agent/actions.py | 377 +++++++++++++++++- src/primaite/game/agent/interface.py | 46 ++- src/primaite/game/agent/observations.py | 243 +++++++----- src/primaite/game/agent/rewards.py | 25 +- src/primaite/game/session.py | 46 +++ 7 files changed, 1382 insertions(+), 174 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 79cfccac..8cf401cc 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -27,11 +27,11 @@ game_config: - type: LOGON - type: LOGOFF applications: - - application_ref: client_2_web_browser - actions: - - type: EXECUTE - execution_definition: - target_address: arcd.com + # - application_ref: client_2_web_browser + # actions: + # - type: EXECUTE + # execution_definition: + # target_address: arcd.com reward_function: null agent_settings: start_step: 5 @@ -42,7 +42,8 @@ game_config: team: RED type: RedDatabaseCorruptingAgent observation_space: - network: + type: UC2RedObservation + options: nodes: - node_ref: client_1 observations: @@ -85,13 +86,307 @@ game_config: - ref: defender team: blue - type: GATE_RL_AGENT + type: GATERLAgent observation_space: - network: + type: UC2BlueObservation + options: nodes: - node_ref: router_1 #TODO: more sub-options here - node_ref: switch_1 - node_ref: switch_2 + - node_ref: domain_controller + services: + - service_ref: domain_controller_dns_server + - node_ref: web_server + services: + - service_ref: web_server_database_client + - node_ref: database_server + services: + - service_ref: database_service + folders: + - folder_name: database + files: + - file_name: database.db + - node_ref: backup_server + # services: + # - service_ref: backup_service + - node_ref: security_suite + - node_ref: client_1 + - node_ref: client_2 + links: + - link_ref: router_1___switch_1 + - link_ref: router_1___switch_2 + - link_ref: switch_1___domain_controller + - link_ref: switch_1___web_server + - link_ref: switch_1___database_server + - link_ref: switch_1___backup_server + - link_ref: switch_1___security_suite + - link_ref: switch_2___client_1 + - link_ref: switch_2___client_2 + - link_ref: switch_2___security_suite + acl: + router_node_ref: router_1 + ics: null + + + action_space: + action_list: + - DONOTHING + - NODE_SERVICE_SCAN + - NODE_SERVICE_STOP + # - NODE_SERVICE_START + # - NODE_SERVICE_PAUSE + # - NODE_SERVICE_RESUME + # - NODE_SERVICE_RESTART + # - NODE_SERVICE_DISABLE + # - NODE_SERVICE_ENABLE + # - NODE_FILE_SCAN + # - NODE_FILE_CHECKHASH + # - NODE_FILE_DELETE + # - NODE_FILE_REPAIR + # - NODE_FILE_RESTORE + # - NODE_FOLDER_SCAN + # - NODE_FOLDER_CHECKHASH + # - NODE_FOLDER_REPAIR + # - NODE_FOLDER_RESTORE + # - NODE_OS_SCAN + # - NODE_SHUTDOWN + # - NODE_STARTUP + # - NODE_RESET + # - NETWORK_ACL_ADDRULE + # - NETWORK_ACL_REMOVERULE + # - NETWORK_NIC_ENABLE + - NETWORK_NIC_DISABLE + + action_map: + 0: + - action: DONOTHING + options: {} + # scan webapp service + 1: + - action: NODE_SERVICE_SCAN + options: + - node_id: 2 + - service_id: 1 + # stop webapp service + 2: + - action: NODE_SERVICE_STOP + options: + - node_id: 2 + - service_id: 1 + # start webapp service + 3: + - action: "NODE_SERVICE_START" + options: + - node_id: 2 + - service_id: 1 + 4: + - action: "NODE_SERVICE_PAUSE" + options: + - node_id: 2 + - service_id: 1 + 5: + - action: "NODE_SERVICE_RESUME" + options: + - node_id: 2 + - service_id: 1 + 6: + - action: "NODE_SERVICE_RESTART" + options: + - node_id: 2 + - service_id: 1 + 7: + - action: "NODE_SERVICE_DISABLE" + options: + - node_id: 2 + - service_id: 1 + 8: + - action: "NODE_SERVICE_ENABLE" + options: + - node_id: 2 + - service_id: 1 + 9: + - action: "NODE_FILE_SCAN" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 + 10: + - action: "NODE_FILE_CHECKHASH" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 + 11: + - action: "NODE_FILE_DELETE" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 + 12: + - action: "NODE_FILE_REPAIR" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 + 13: + - action: "NODE_FILE_RESTORE" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 + 14: + - action: "NODE_FOLDER_SCAN" + options: + - node_id: 3 + - folder_id: 1 + 15: + - action: "NODE_FOLDER_CHECKHASH" + options: + - node_id: 3 + - folder_id: 1 + 16: + - action: "NODE_FOLDER_REPAIR" + options: + - node_id: 3 + - folder_id: 1 + 17: + - action: "NODE_FOLDER_RESTORE" + options: + - node_id: 3 + - folder_id: 1 + 18: + - action: "NODE_OS_SCAN" + options: + - node_id: 3 + 19: + - action: "NODE_SHUTDOWN" + options: + - node_id: 6 + 20: + - action: "NODE_STARTUP" + options: + - node_id: 6 + 21: + - action: "NODE_RESET" + options: + - node_id: 6 + 22: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 6 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 23: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 5 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 24: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 4 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 25: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 3 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 26: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 2 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 27: + - action: "NETWORK_ACL_ADDRULE" + options: + - position: 1 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... + 28: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 0 + 29: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 1 + 30: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 2 + 31: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 3 + 32: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 4 + 33: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 5 + 34: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 6 + 35: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 7 + 36: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 8 + 37: + - action: "NETWORK_ACL_REMOVERULE" + options: + - position: 9 + 38: + - action: "NETWORK_NIC_DISABLE" + options: + - node_id: 6 + - nic_index: 1 + 39: + - action: "NETWORK_NIC_ENABLE" + options: + - node_id: 6 + - nic_index: 1 + + options: + nodes: + - node_ref: router_1 + - node_ref: switch_1 + - node_ref: switch_2 - node_ref: domain_controller - node_ref: web_server - node_ref: database_server @@ -99,18 +394,13 @@ game_config: - node_ref: security_suite - node_ref: client_1 - node_ref: client_2 - links: - - link_ref: ... # - acl: ... # - ics: ... # + max_folders_per_node: 2 + max_files_per_folder: 2 + max_services_per_node: 2 + max_nics_per_node: 8 + max_acl_rules: 10 - action_space: - actions: - - type: DO_NOTHING - network: - nodes: - - node_ref: router_1 reward_function: # ... agent_settings: @@ -175,7 +465,6 @@ simulation: domain_mapping: arcd.com: 192.168.1.12 # web server - - ref: web_server type: server hostname: web_server @@ -200,7 +489,6 @@ simulation: - ref: database_service type: DatabaseService - - ref: backup_server type: server hostname: backup_server @@ -224,7 +512,6 @@ simulation: ip_address: 192.168.10.110 subnet_mask: 255.255.255.0 - - ref: client_1 type: computer hostname: client_1 @@ -251,7 +538,6 @@ simulation: - ref: client_2_dns_client type: DNSClient - links: - ref: router_1___switch_1 endpoint_a_ref: router_1 diff --git a/sandbox.ipynb b/sandbox.ipynb index 05efcfa2..3ff72170 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -2,9 +2,18 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 13, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The autoreload extension is already loaded. To reload it, use:\n", + " %reload_ext autoreload\n" + ] + } + ], "source": [ "%load_ext autoreload\n", "%autoreload 2" @@ -12,7 +21,381 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.game.session import PrimaiteSession\n", + "from primaite.simulator.sim_container import Simulation\n", + "from primaite.game.agent.interface import AbstractAgent\n", + "from primaite.simulator.network.networks import arcd_uc2_network\n" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "sess = PrimaiteSession()" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [], + "source": [ + "network = sess.simulation.network" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [], + "source": [ + "from ipaddress import IPv4Address\n", + "\n", + "from primaite.simulator.network.container import Network\n", + "from primaite.simulator.network.hardware.base import NIC\n", + "from primaite.simulator.network.hardware.nodes.computer import Computer\n", + "from primaite.simulator.network.hardware.nodes.router import ACLAction, Router\n", + "from primaite.simulator.network.hardware.nodes.server import Server\n", + "from primaite.simulator.network.hardware.nodes.switch import Switch\n", + "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", + "from primaite.simulator.network.transmission.transport_layer import Port\n", + "from primaite.simulator.system.applications.database_client import DatabaseClient\n", + "from primaite.simulator.system.services.database_service import DatabaseService\n", + "from primaite.simulator.system.services.dns_client import DNSClient\n", + "from primaite.simulator.system.services.dns_server import DNSServer\n", + "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-02 15:10:20,422: Added node 6abb7664-4d17-45ff-a3c7-dbcccffcfd6d to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,424: Added node 3edbc521-3c80-47e3-8017-dbc38fb00a73 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,428: Added node 94457fb9-04a1-4dc1-9ff7-b64df0da7424 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,432: Added node 0d311d72-139c-41bf-aef7-fa9b01b124d7 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,439: Added node 6161e785-f377-48de-aa4f-20d3646da635 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,444: Added node 55a9e9f8-ee3a-4c28-9b6d-c0c0d78a3f6a to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,447: Added node 2f04ca45-3439-489a-81f7-41cca5ae8adc to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,531: Added node 98660c30-8e48-4b96-967a-d62ca71b4d6d to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,545: Added node 1a184184-b204-40de-986a-4d7459036dbe to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,551: Added node 17b92b9a-6805-4677-85f7-c0c0521a6e25 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", + "2023-10-02 15:10:20,555::ERROR::primaite.simulator.network.hardware.base::175::NIC da:f3:1b:87:24:20/192.168.10.110 cannot be enabled as it is not connected to a Link\n" + ] + } + ], + "source": [ + "router_1 = Router(hostname=\"router_1\", num_ports=5)\n", + "router_1.power_on()\n", + "router_1.configure_port(port=1, ip_address=\"192.168.1.1\", subnet_mask=\"255.255.255.0\")\n", + "router_1.configure_port(port=2, ip_address=\"192.168.10.1\", subnet_mask=\"255.255.255.0\")\n", + "\n", + "# Switch 1\n", + "switch_1 = Switch(hostname=\"switch_1\", num_ports=8)\n", + "switch_1.power_on()\n", + "network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[8])\n", + "router_1.enable_port(1)\n", + "\n", + "# Switch 2\n", + "switch_2 = Switch(hostname=\"switch_2\", num_ports=8)\n", + "switch_2.power_on()\n", + "network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[8])\n", + "router_1.enable_port(2)\n", + "\n", + "# Client 1\n", + "client_1 = Computer(\n", + " hostname=\"client_1\",\n", + " ip_address=\"192.168.10.21\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.10.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "client_1.power_on()\n", + "client_1.software_manager.install(DNSClient)\n", + "client_1_dns_client_service: DNSServer = client_1.software_manager.software[\"DNSClient\"] # noqa\n", + "client_1_dns_client_service.start()\n", + "network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])\n", + "client_1.software_manager.install(DataManipulationBot)\n", + "db_manipulation_bot: DataManipulationBot = client_1.software_manager.software[\"DataManipulationBot\"]\n", + "db_manipulation_bot.configure(server_ip_address=IPv4Address(\"192.168.1.14\"), payload=\"DROP TABLE IF EXISTS user;\")\n", + "\n", + "# Client 2\n", + "client_2 = Computer(\n", + " hostname=\"client_2\",\n", + " ip_address=\"192.168.10.22\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.10.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "client_2.power_on()\n", + "client_2.software_manager.install(DNSClient)\n", + "client_2_dns_client_service: DNSServer = client_2.software_manager.software[\"DNSClient\"] # noqa\n", + "client_2_dns_client_service.start()\n", + "network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2])\n", + "\n", + "# Domain Controller\n", + "domain_controller = Server(\n", + " hostname=\"domain_controller\",\n", + " ip_address=\"192.168.1.10\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.1.1\",\n", + ")\n", + "domain_controller.power_on()\n", + "domain_controller.software_manager.install(DNSServer)\n", + "\n", + "network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1])\n", + "\n", + "# Database Server\n", + "database_server = Server(\n", + " hostname=\"database_server\",\n", + " ip_address=\"192.168.1.14\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.1.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "database_server.power_on()\n", + "network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3])\n", + "\n", + "ddl = \"\"\"\n", + "CREATE TABLE IF NOT EXISTS user (\n", + "id INTEGER PRIMARY KEY AUTOINCREMENT,\n", + "name VARCHAR(50) NOT NULL,\n", + "email VARCHAR(50) NOT NULL,\n", + "age INT,\n", + "city VARCHAR(50),\n", + "occupation VARCHAR(50)\n", + ");\"\"\"\n", + "\n", + "user_insert_statements = [\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('John Doe', 'johndoe@example.com', 32, 'New York', 'Engineer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Jane Smith', 'janesmith@example.com', 27, 'Los Angeles', 'Designer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Bob Johnson', 'bobjohnson@example.com', 45, 'Chicago', 'Manager');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Alice Lee', 'alicelee@example.com', 22, 'San Francisco', 'Student');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('David Kim', 'davidkim@example.com', 38, 'Houston', 'Consultant');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Emily Chen', 'emilychen@example.com', 29, 'Seattle', 'Software Developer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Frank Wang', 'frankwang@example.com', 55, 'New York', 'Entrepreneur');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Grace Park', 'gracepark@example.com', 31, 'Los Angeles', 'Marketing Specialist');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Henry Wu', 'henrywu@example.com', 40, 'Chicago', 'Accountant');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Isabella Kim', 'isabellakim@example.com', 26, 'San Francisco', 'Graphic Designer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Jake Lee', 'jakelee@example.com', 33, 'Houston', 'Sales Manager');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Kelly Chen', 'kellychen@example.com', 28, 'Seattle', 'Web Developer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Lucas Liu', 'lucasliu@example.com', 42, 'New York', 'Lawyer');\", # noqa\n", + " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Maggie Wang', 'maggiewang@example.com', 30, 'Los Angeles', 'Data Analyst');\", # noqa\n", + "]\n", + "database_server.software_manager.install(DatabaseService)\n", + "database_service: DatabaseService = database_server.software_manager.software[\"DatabaseService\"] # noqa\n", + "database_service.start()\n", + "database_service._process_sql(ddl, None) # noqa\n", + "for insert_statement in user_insert_statements:\n", + " database_service._process_sql(insert_statement, None) # noqa\n", + "\n", + "# Web Server\n", + "web_server = Server(\n", + " hostname=\"web_server\",\n", + " ip_address=\"192.168.1.12\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.1.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "web_server.power_on()\n", + "web_server.software_manager.install(DatabaseClient)\n", + "\n", + "database_client: DatabaseClient = web_server.software_manager.software[\"DatabaseClient\"]\n", + "database_client.configure(server_ip_address=IPv4Address(\"192.168.1.14\"))\n", + "network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])\n", + "database_client.run()\n", + "database_client.connect()\n", + "\n", + "# register the web_server to a domain\n", + "dns_server_service: DNSServer = domain_controller.software_manager.software[\"DNSServer\"] # noqa\n", + "dns_server_service.start()\n", + "dns_server_service.dns_register(\"arcd.com\", web_server.ip_address)\n", + "\n", + "# Backup Server\n", + "backup_server = Server(\n", + " hostname=\"backup_server\",\n", + " ip_address=\"192.168.1.16\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.1.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "backup_server.power_on()\n", + "network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4])\n", + "\n", + "# Security Suite\n", + "security_suite = Server(\n", + " hostname=\"security_suite\",\n", + " ip_address=\"192.168.1.110\",\n", + " subnet_mask=\"255.255.255.0\",\n", + " default_gateway=\"192.168.1.1\",\n", + " dns_server=IPv4Address(\"192.168.1.10\"),\n", + ")\n", + "security_suite.power_on()\n", + "network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7])\n", + "security_suite.connect_nic(NIC(ip_address=\"192.168.10.110\", subnet_mask=\"255.255.255.0\"))\n", + "network.connect(endpoint_b=security_suite.ethernet_port[2], endpoint_a=switch_2.switch_ports[7])\n", + "\n", + "router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)\n", + "\n", + "router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)\n", + "\n", + "# Allow PostgreSQL requests\n", + "router_1.acl.add_rule(\n", + " action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0\n", + ")\n", + "\n", + "# Allow DNS requests\n", + "router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [], + "source": [ + "node_uuid_list = list(sess.simulation.network.nodes.keys())" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.game.agent.actions import ActionManager" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [], + "source": [ + "actman = ActionManager(sess.simulation, [\"DONOTHING\", \"NODE_SERVICE_SCAN\", \"NODE_SERVICE_STOP\", \"NODE_FOLDER_SCAN\"],node_uuid_list,act_map={\n", + " 0:{\n", + " \"action\": \"DONOTHING\",\n", + " \"options\": {}\n", + " },\n", + " 1:{\n", + " \"action\": \"NODE_SERVICE_SCAN\",\n", + " \"options\": {\"node_id\":0, \"service_id\":0},\n", + " },\n", + " 2:{\n", + " \"action\": \"NODE_SERVICE_SCAN\",\n", + " \"options\": {\"node_id\":1, \"service_id\":0},\n", + " },\n", + " 3:{\n", + " \"action\": \"NODE_FOLDER_SCAN\",\n", + " \"options\": {\"node_id\":4, \"folder_id\":0},\n", + " }\n", + "})" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [], + "source": [ + "act_id, act_options = actman.get_action(3)\n", + "my_trial_act = actman.form_request(action_identifier=act_id, action_options=act_options)" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [], + "source": [ + "sess.simulation.apply_action(my_trial_act)" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['network',\n", + " 'node',\n", + " '6161e785-f377-48de-aa4f-20d3646da635',\n", + " 'file_system',\n", + " 'folder',\n", + " '5aefe92b-923c-4684-b3bf-e78dd18d4771',\n", + " 'scan']" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "my_trial_act" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sess.step()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sess.step_counter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from gym import spaces" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sp = spaces.Tuple( (spaces.MultiDiscrete([3, 2]), spaces.MultiDiscrete([3, 2]), spaces.MultiDiscrete([3, 2]),))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sp.sample()" + ] + }, + { + "cell_type": "code", + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -39,40 +422,16 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-09-26 12:19:50,895: Added node 0fb262e1-a714-420a-aec7-be37f0deeb75 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,898: Added node 310ca8d7-01e0-401e-b604-705c290e5376 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,900: Added node b3b08f1f-7805-47b2-bdb6-3d83098cd740 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,903: Added node adb37f3e-2307-4123-bff3-01f125883be8 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,906: Added node 1a490716-2ccd-452d-b87e-324d29120b59 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,911: Added node 033460d8-0249-4bdd-aaf0-751b24cc0a1e to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,914: Added node 1e7e4e49-78bf-4031-8372-ee71902720f3 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,916: Added node c9f24a13-e5c8-437b-9234-b0c3f8120e2c to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,920: Added node c881f3ee-2176-493b-a6c2-cad829bf0b6d to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", - "2023-09-26 12:19:50,922: Added node a3ea75d8-bc2c-4713-92a4-2588b4f43ed6 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "service type not found DatabaseBackup\n", - "service type not found WebBrowser\n" - ] - } - ], + "outputs": [], "source": [ "# import yaml\n", "\n", "\n", "from typing import Dict\n", "from primaite.game.agent.interface import AbstractAgent\n", + "from primaite.game.agent.observations import AclObservation, FileObservation, FolderObservation, ICSObservation, LinkObservation, NicObservation, NodeObservation, NullObservation, ServiceObservation, UC2BlueObservation, UC2RedObservation\n", "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", "from primaite.simulator.system.services.service import Service\n", "\n", @@ -130,8 +489,8 @@ " subnet_mask=port_cfg['subnet_mask'])\n", " if 'acl' in node_cfg:\n", " for r_num, r_cfg in node_cfg['acl'].items():\n", - " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, so that we can do\n", - " # both of these things once: check if a key is defined, access and convert it to a \n", + " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating \n", + " # this: 'r_cfg.get('src_port')'\n", " # Port/IPProtocol. TODO Refactor\n", " new_node.acl.add_rule(\n", " action = ACLAction[r_cfg['action']],\n", @@ -211,8 +570,70 @@ " action_space_cfg = agent_cfg['action_space']\n", " observation_space_cfg = agent_cfg['observation_space']\n", " reward_function_cfg = agent_cfg['reward_function']\n", + " \n", + " # CREATE OBSERVATION SPACE\n", + " if observation_space_cfg is None:\n", + " obs_space = NullObservation()\n", + " elif observation_space_cfg['type'] == 'UC2BlueObservation':\n", + " node_obs_list = []\n", + " link_obs_list = []\n", + " \n", + " \n", + " #node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses\n", + " node_ip_to_index = {}\n", + " for node_idx, node_cfg in enumerate(nodes_cfg):\n", + " n_ref = node_cfg['ref']\n", + " n_obj = net.nodes[ref_map_nodes[n_ref]]\n", + " for nic_uuid, nic_obj in n_obj.nics.items():\n", + " node_ip_to_index[nic_obj.ip_address] = node_idx + 2\n", + "\n", + " \n", + " \n", + " for node_obs_cfg in observation_space_cfg['options']['nodes']:\n", + " node_ref = node_obs_cfg['node_ref']\n", + " folder_obs_list = []\n", + " service_obs_list = []\n", + " if 'services' in node_obs_cfg:\n", + " for service_obs_cfg in node_obs_cfg['services']:\n", + " service_obs_list.append(ServiceObservation(where=['network','nodes',ref_map_nodes[node_ref],'services',ref_map_services[service_obs_cfg['service_ref']]]))\n", + " if 'folders' in node_obs_cfg:\n", + " for folder_obs_cfg in node_obs_cfg['folders']:\n", + " file_obs_list = []\n", + " if 'files' in folder_obs_cfg:\n", + " for file_obs_cfg in folder_obs_cfg['files']:\n", + " file_obs_list.append(FileObservation(where=['network','nodes',ref_map_nodes[node_ref], 'folders',folder_obs_cfg['folder_name'], 'files', file_obs_cfg['file_name']]))\n", + " folder_obs_list.append(FolderObservation(where=['network','nodes',ref_map_nodes[node_ref], 'folders',folder_obs_cfg['folder_name']], files=file_obs_list))\n", + " nic_obs_list = []\n", + " for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg['node_ref']]].nics.keys():\n", + " nic_obs_list.append(NicObservation(where=['network','nodes',ref_map_nodes[node_ref],'NICs',nic_uuid]))\n", + " node_obs_list.append(NodeObservation(where=['network','nodes',ref_map_nodes[node_ref]], services=service_obs_list, folders=folder_obs_list,nics=nic_obs_list, logon_status=False))\n", + " for link_obs_cfg in observation_space_cfg['options']['links']:\n", + " link_ref = link_obs_cfg['link_ref']\n", + " link_obs_list.append(LinkObservation(where=['network' ,'links', ref_map_links[link_ref]]))\n", + "\n", + " acl_obs = AclObservation(node_ip_to_id=node_ip_to_index, ports=game_cfg['ports'], protocols=game_cfg['ports'], where=['network','nodes',observation_space_cfg['options']['acl']['router_node_ref']])\n", + " obs_space = UC2BlueObservation(nodes=node_obs_list,links=link_obs_list,acl=acl_obs, ics=ICSObservation())\n", + " elif observation_space_cfg['type'] == 'UC2RedObservation':\n", + " obs_space = UC2RedObservation.from_config(observation_space_cfg['options'], sim=sim)\n", + " else:\n", + " print(\"observation space config not specified correctly.\")\n", + " obs_space = NullObservation()\n", + " \n", + " # CREATE ACTION SPACE\n", + " \n", + "\n", + "\n", + " # CREATE REWARD FUNCTION\n", + "\n", + " # CREATE AGENT\n", " if agent_type == 'GreenWebBrowsingAgent':\n", - " new_agent = GreenWebBrowsingAgent()\n", + " ...\n", + " elif agent_type == 'GATERLAgent':\n", + " ...\n", + " elif agent_type == 'RedDatabaseCorruptingAgent':\n", + " ...\n", + " else:\n", + " print(\"agent type not found\")\n", "\n", "\n", " #4. set up agents' actions and observation spaces.\n", @@ -228,7 +649,7 @@ "metadata": {}, "outputs": [], "source": [ - "print(s.simulation.describe_state())" + "s.agents" ] }, { diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index cefd9917..cb7061fc 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1,21 +1,374 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, List - -from pydantic import BaseModel +from typing import Any, Dict, List, Optional, Tuple +import itertools -class AbstractAction(BaseModel): +from primaite.simulator.sim_container import Simulation + +from gym import spaces + +class ExecutionDefiniton(ABC): + """ + Converter from actions to simulator requests. + + Allows adding extra data/context that defines in more detail what an action means. + """ + + """ + Examples: + ('node', 'service', 'scan', 2, 0) means scan the first service on node index 2 + -> ['network', 'nodes', , 'services', , 'scan'w] + """ + ... + + +class AbstractAction(ABC): + @abstractmethod - def __call__(self, action: Any) -> List[str]: - """_summary_ - - :param action: _description_ - :type action: Any - :return: _description_ - :rtype: List[str] + def __init__(self, manager:"ActionManager", **kwargs) -> None: """ + Init method for action. + + All action init functions should accept **kwargs as a way of ignoring extra arguments. + + Since many parameters are defined for the action space as a whole (such as max files per folder, max services + per node), we need to pass those options to every action that gets created. To pervent verbosity, these + parameters are just broadcasted to all actions and the actions can pay attention to the ones that apply. + """ + self.name:str = "" + """Human-readable action identifier used for printing, logging, and reporting.""" + self.shape = (0,) + """Tuple describing number of options for each parameter of this action. Can be passed to + gym.spaces.MultiDiscrete to form a valid space.""" + self.manager:ActionManager = manager + + + @abstractmethod + def form_request(self) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [] + + +class DoNothingAction(AbstractAction): + def __init__(self, manager:"ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + self.name = "DONOTHING" + self.shape = (1,) + + def form_request(self) -> List[str]: + return ["do_nothing"] + +class NodeServiceAbstractAction(AbstractAction): + """ + Base class for service actions. + + Any action which applies to a service and uses node_id and service_id as its only two parameters can inherit from + this base class. + """ + @abstractmethod + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Tuple[int] = (num_nodes, num_services) + self.verb:str + + def form_request(self, node_id:int, service_id:int) -> List[str]: + node_uuid = self.manager.get_node_uuid_by_idx(node_id) + service_uuid = self.manager.get_service_uuid_by_idx(node_id, service_id) + if node_uuid is None or service_uuid is None: + return ["do_nothing"] + return ['network', 'node', node_uuid, 'services', service_uuid, self.verb] + +class NodeServiceScanAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "scan" + +class NodeServiceStopAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "stop" + +class NodeServiceStartAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "start" + +class NodeServicePauseAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "pause" + +class NodeServiceResumeAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "resume" + +class NodeServiceRestartAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "restart" + +class NodeServiceDisableAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "disable" + +class NodeServiceEnableAction(NodeServiceAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = "enable" + + + +class NodeFolderAbstractAction(AbstractAction): + @abstractmethod + def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + super().__init__(manager=manager) + self.shape = (num_nodes, num_folders) + self.verb: str + + def form_request(self, node_id:int, folder_id:int) -> List[str]: + node_uuid = self.manager.get_node_uuid_by_idx(node_id) + folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) + if node_uuid is None or folder_uuid is None: + return ["do_nothing"] + return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, self.verb] + +class NodeFolderScanAction(NodeFolderAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, **kwargs) + self.verb:str = "scan" + +class NodeFolderCheckhashAction(NodeFolderAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, **kwargs) + self.verb:str = "checkhash" + +class NodeFolderRepairAction(NodeFolderAbstractAction): + def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, **kwargs) + self.verb:str = "repair" + +class NodeFolderRestoreAction(NodeFolderAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, **kwargs) + self.verb:str = "restore" + + +class NodeFileAbstractAction(AbstractAction): + @abstractmethod + def __init__(self, manager:"ActionManager", num_nodes:int, num_folders:int, num_files:int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape:Tuple[int] = (num_nodes, num_folders, num_files) + self.verb:str + + def form_request(self, node_id:int, folder_id:int, file_id:int) -> List[str]: + node_uuid = self.manager.get_node_uuid_by_idx(node_id) + folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) + file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id) + if node_uuid is None or folder_uuid is None or file_uuid is None: + return ["do_nothing"] + return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, 'files', file_uuid, self.verb] + +class NodeFileScanAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + self.verb = "scan" + +class NodeFileCheckhashAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + self.verb = "checkhash" + +class NodeFileDeleteAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + self.verb = "delete" + +class NodeFileRepairAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + self.verb = "repair" + +class NodeFileRestoreAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + self.verb = "restore" + +class NodeAbstractAction(AbstractAction): + @abstractmethod + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Tuple[int] = (num_nodes,) + self.verb: str + + def form_request(self, node_id:int) -> List[str]: + node_uuid = self.manager.get_node_uuid_by_idx(node_id) + return ["network", "node", node_uuid, self.verb] + +class NodeOSScanAction(NodeAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = 'scan' + +class NodeShutdownAction(NodeAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = 'shutdown' + +class NodeStartupAction(NodeAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = 'start' + +class NodeResetAction(NodeAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.verb = 'reset' + +class NetworkACLAddRuleAction(AbstractAction): + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + num_permissions = 2 + self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_nics, num_nics, num_ports, num_ports, num_protocols) + + + + + + +class ActionManager: + # let the action manager handle the conversion of action spaces into a single discrete integer space. + # + + + # when action space is created, it will take subspaces and generate an action map by enumerating all possibilities, + # BUT, the action map can be provided in the config, in which case it will use that. + + # action map is basically just a mapping between integer and CAOS action (incl. parameter values) + # for example the action map can be: + # 0: DONOTHING + # 1: NODE, FILE, SCAN, NODEID=2, FOLDERID=1, FILEID=0 + # 2: ...... + __act_class_identifiers:Dict[str,type] = { + "DONOTHING": DoNothingAction, + "NODE_SERVICE_SCAN": NodeServiceScanAction, + "NODE_SERVICE_STOP": NodeServiceStopAction, + # "NODE_SERVICE_START": NodeServiceStartAction, + # "NODE_SERVICE_PAUSE": NodeServicePauseAction, + # "NODE_SERVICE_RESUME": NodeServiceResumeAction, + # "NODE_SERVICE_RESTART": NodeServiceRestartAction, + # "NODE_SERVICE_DISABLE": NodeServiceDisableAction, + # "NODE_SERVICE_ENABLE": NodeServiceEnableAction, + # "NODE_FILE_SCAN": NodeFileScanAction, + # "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, + # "NODE_FILE_DELETE": NodeFileDeleteAction, + # "NODE_FILE_REPAIR": NodeFileRepairAction, + # "NODE_FILE_RESTORE": NodeFileRestoreAction, + "NODE_FOLDER_SCAN": NodeFolderScanAction, + # "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, + # "NODE_FOLDER_REPAIR": NodeFolderRepairAction, + # "NODE_FOLDER_RESTORE": NodeFolderRestoreAction, + # "NODE_OS_SCAN": NodeOSScanAction, + # "NODE_SHUTDOWN": NodeShutdownAction, + # "NODE_STARTUP": NodeStartupAction, + # "NODE_RESET": NodeResetAction, + # "NETWORK_ACL_ADDRULE": NetworkACLAddRuleAction, + # "NETWORK_ACL_REMOVERULE": NetworkACLRemoveRuleAction, + # "NETWORK_NIC_ENABLE": NetworkNICEnable, + # "NETWORK_NIC_DISABLE": NetworkNICDisable, + } + + + def __init__(self, + sim:Simulation, + actions:List[str], + node_uuids:List[str], + max_folders_per_node:int = 2, + max_files_per_folder:int = 2, + max_services_per_node:int = 2, + max_nics_per_node:int=8, + max_acl_rules:int=10, + act_map:Optional[Dict[int, Dict]]=None) -> None: + self.sim: Simulation = sim + self.node_uuids:List[str] = node_uuids + + action_args = { + "num_nodes": len(node_uuids), + "num_folders":max_folders_per_node, + "num_files": max_files_per_folder, + "num_services": max_services_per_node, + "num_nics": max_nics_per_node, + "num_acl_rules": max_acl_rules} + self.actions: Dict[str, AbstractAction] = {} + for act_type in actions: + self.actions[act_type] = self.__act_class_identifiers[act_type](self, **action_args) + + self.action_map:Dict[int, Tuple[str, Dict]] = {} + """ + Action mapping that converts an integer to a specific action and parameter choice. + + For example : + {0: ("NODE_SERVICE_SCAN", {node_id:0, service_id:2})} + """ + if act_map is None: + self.action_map = self._enumerate_actions() + else: + self.action_map = {i:(a['action'], a['options']) for i,a in act_map.items()} + # make sure all numbers between 0 and N are represented as dict keys in action map + assert all([i in self.action_map.keys() for i in range(len(self.action_map))]) + + def _enumerate_actions(self,) -> Dict[int, Tuple[AbstractAction, Dict]]: ... + def get_action(self, action: int) -> Tuple[str,Dict]: + """Produce action in CAOS format""" + """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" + """The caos format is basically a action identifier, followed by parameters stored in a dictionary""" + act_identifier, act_options = self.action_map[action] + return act_identifier, act_options -class ActionSpace: + def form_request(self, action_identifier:str, action_options:Dict): + """Take action in CAOS format and use the execution definition to change it into PrimAITE request format""" + act_obj = self.actions[action_identifier] + return act_obj.form_request(**action_options) + + @property + def space(self) -> spaces.Space: + return spaces.Discrete(len(self.action_map)) + + def get_node_uuid_by_idx(self, node_idx): + return self.node_uuids[node_idx] + + def get_folder_uuid_by_idx(self, node_idx, folder_idx) -> Optional[str]: + node_uuid = self.get_node_uuid_by_idx(node_idx) + node = self.sim.network.nodes[node_uuid] + folder_uuids = list(node.file_system.folders.keys()) + return folder_uuids[folder_idx] if len(folder_uuids)>folder_idx else None + + def get_file_uuid_by_idx(self, node_idx, folder_idx, file_idx) -> Optional[str]: + node_uuid = self.get_node_uuid_by_idx(node_idx) + node = self.sim.network.nodes[node_uuid] + folder_uuids = list(node.file_system.folders.keys()) + if len(folder_uuids)<=folder_idx: + return None + folder = node.file_system.folders[folder_uuids[folder_idx]] + file_uuids = list(folder.files.keys()) + return file_uuids[file_idx] if len(file_uuids)>file_idx else None + + def get_service_uuid_by_idx(self, node_idx, service_idx) -> Optional[str]: + node_uuid = self.get_node_uuid_by_idx(node_idx) + node = self.sim.network.nodes[node_uuid] + service_uuids = list(node.services.keys()) + return service_uuids[service_idx] if len(service_uuids)>service_idx else None + + + + + + +class UC2RedActions(AbstractAction): + ... + +class UC2GreenActionSpace(ActionManager): ... diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index b1ade94b..0e682b60 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -2,32 +2,70 @@ # That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word # If you disagree, make a comment in the PR review and we can discuss from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union, TypeAlias +import numpy as np -from primaite.game.agent.actions import ActionSpace +from primaite.game.agent.actions import ActionManager from primaite.game.agent.observations import ObservationSpace from primaite.game.agent.rewards import RewardFunction +ObsType:TypeAlias = Union[Dict, np.ndarray] class AbstractAgent(ABC): """Base class for scripted and RL agents.""" def __init__( self, - action_space: Optional[ActionSpace], + action_space: Optional[ActionManager], observation_space: Optional[ObservationSpace], reward_function: Optional[RewardFunction], ) -> None: - self.action_space: Optional[ActionSpace] = action_space + self.action_space: Optional[ActionManager] = action_space self.observation_space: Optional[ObservationSpace] = observation_space self.reward_function: Optional[RewardFunction] = reward_function + # exection definiton converts CAOS action to Primaite simulator request, sometimes having to enrich the info + # by for example specifying target ip addresses, or converting a node ID into a uuid + self.execution_definition = None + + def get_obs_from_state(self, state:Dict) -> ObsType: + """ + state : dict state directly from simulation.describe_state + output : dict state according to CAOS. + """ + return self.observation_space.observe(state) + + def get_reward_from_state(self, state:Dict) -> float: + return self.reward_function.calculate(state) + + @abstractmethod + def get_action(self, obs:ObsType, reward:float=None): + # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 1-40, + # then use a bespoke conversion to take 1-40 int back into CAOS action + return ('NODE', 'SERVICE', 'SCAN', '', '') + + @abstractmethod + def format_request(self, action) -> List[str]: + # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. + # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR + """Format action into format expected by the simulator, and apply execution definition if applicable.""" + return ['network', 'nodes', '', 'file_system', 'folder', 'root', 'scan'] + + + + class AbstractScriptedAgent(AbstractAgent): """Base class for actors which generate their own behaviour.""" ... +class RandomAgent(AbstractScriptedAgent): + """Agent that ignores its observation and acts completely at random.""" + + def get_action(self, obs:ObsType, reward:float=None): + return self.action_space.space.sample() + class AbstractGATEAgent(AbstractAgent): """Base class for actors controlled via external messages, such as RL policies.""" diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 4d4796e1..f919a723 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -55,7 +55,7 @@ class AbstractObservation(ABC): class FileObservation(AbstractObservation): - def __init__(self, where: List[str] = []) -> None: + def __init__(self, where: Optional[List[str]] = None) -> None: """ _summary_ @@ -68,12 +68,12 @@ class FileObservation(AbstractObservation): :type where: Optional[List[str]] """ super().__init__() - self.where: List[str] = where + self.where: Optional[List[str]] = where self.default_observation: spaces.Space = {"health_status": 0} "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted." def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation file_state = access_from_nested_dict(state, self.where) if file_state is NOT_PRESENT_IN_STATE: @@ -89,21 +89,21 @@ class ServiceObservation(AbstractObservation): default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} "Default observation is what should be returned when the service doesn't exist." - def __init__(self, where: List[str] = []) -> None: + def __init__(self, where: Optional[List[str]] = None) -> None: """ :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with zeroes. A typical location for a service looks like this: - `['network','nodes',,'servics', ]` + `['network','nodes',,'services', ]` :type where: Optional[List[str]] """ super().__init__() - self.where: List[str] = where + self.where: Optional[List[str]] = where def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation service_state = access_from_nested_dict(state, self.where) @@ -120,7 +120,7 @@ class LinkObservation(AbstractObservation): default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} "Default observation is what should be returned when the link doesn't exist." - def __init__(self, where: List[str] = []) -> None: + def __init__(self, where: Optional[List[str]] = None) -> None: """ :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with @@ -131,10 +131,10 @@ class LinkObservation(AbstractObservation): :type where: Optional[List[str]] """ super().__init__() - self.where: List[str] = where + self.where: Optional[List[str]] = where def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation link_state = access_from_nested_dict(state, self.where) @@ -156,7 +156,7 @@ class LinkObservation(AbstractObservation): class FolderObservation(AbstractObservation): - def __init__(self, where: List[str] = [], files: List[FileObservation] = []) -> None: + def __init__(self, where: Optional[List[str]] = None, files: List[FileObservation] = []) -> None: """Initialise folder Observation, including files inside of the folder. :param where: Where in the simulation state dictionary to find the relevant information for this folder. @@ -175,7 +175,7 @@ class FolderObservation(AbstractObservation): """ super().__init__() - self.where: List[str] = where + self.where: Optional[List[str]] = where self.files: List[FileObservation] = files @@ -185,7 +185,7 @@ class FolderObservation(AbstractObservation): } def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation folder_state = access_from_nested_dict(state, self.where) if folder_state is NOT_PRESENT_IN_STATE: @@ -213,12 +213,12 @@ class FolderObservation(AbstractObservation): class NicObservation(AbstractObservation): default_observation: spaces.Space = {"nic_status": 0} - def __init__(self, where: List[str] = []) -> None: - super.__init__() - self.where: List[str] = where + def __init__(self, where: Optional[List[str]] = None) -> None: + super().__init__() + self.where: Optional[List[str]] = where def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation nic_state = access_from_nested_dict(state, self.where) if nic_state is NOT_PRESENT_IN_STATE: @@ -234,10 +234,11 @@ class NicObservation(AbstractObservation): class NodeObservation(AbstractObservation): def __init__( self, - where: List[str] = [], + where: Optional[List[str]] = None, services: List[ServiceObservation] = [], folders: List[FolderObservation] = [], nics: List[NicObservation] = [], + logon_status:bool=False ) -> None: """ Configurable observation for a node in the simulation. @@ -259,12 +260,13 @@ class NodeObservation(AbstractObservation): :param max_nics: Max number of NICS in this node's obs space, defaults to 5 :type max_nics: int, optional """ - super.__init__() - self.where: List[str] = where + super().__init__() + self.where: Optional[List[str]] = where self.services: List[ServiceObservation] = services self.folders: List[FolderObservation] = folders self.nics: List[NicObservation] = nics + self.logon_status:bool=logon_status self.default_observation: Dict = { "SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)}, @@ -272,9 +274,11 @@ class NodeObservation(AbstractObservation): "NICS": {i + 1: n.default_observation for i, n in enumerate(self.nics)}, "operating_status": 0, } + if self.logon_status: + self.default_observation['logon_status']=0 def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation node_state = access_from_nested_dict(state, self.where) @@ -288,18 +292,24 @@ class NodeObservation(AbstractObservation): obs["operating_status"] = node_state["operating_state"] obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)} + if self.logon_status: + obs['logon_status'] = 0 + return obs @property def space(self) -> spaces.Space: - return spaces.Dict( - { - "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}), - "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}), - "operating_status": spaces.Discrete(0), - "NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}), - } - ) + space_shape = { + "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}), + "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}), + "operating_status": spaces.Discrete(5), + "NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}), + } + if self.logon_status: + space_shape['logon_status'] = spaces.Discrete(3) + + return spaces.Dict(space_shape) + class AclObservation(AbstractObservation): @@ -308,41 +318,33 @@ class AclObservation(AbstractObservation): # if a file is created at runtime, we have currently got no way of telling the observation space to track it. # this needs adding, but not for the MVP. def __init__( - self, nodes: List[str], ports: List[int], protocols: list[str], where: List[str] = [], num_rules: int = 10 + self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[List[str]] = None, num_rules: int = 10 ) -> None: super().__init__() - self.where: List[str] = where + self.where: Optional[List[str]] = where self.num_rules: int = num_rules - self.node_to_id: Dict[str, int] = {node: i + 1 for i, node in enumerate(nodes)} + self.node_to_id: Dict[str, int] = node_ip_to_id "List of node IP addresses, order in this list determines how they are converted to an ID" - self.port_to_id: Dict[int, int] = {port: i + 1 for i, port in enumerate(ports)} + self.port_to_id: Dict[int, int] = {port: i + 2 for i, port in enumerate(ports)} "List of ports which are part of the game that define the ordering when converting to an ID" - self.protocol_to_id: Dict[str, int] = {protocol: i + 1 for i, protocol in enumerate(protocols)} + self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)} "List of protocols which are part of the game, defines ordering when converting to an ID" - self.default_observation: spaces.Space = spaces.Dict( - { - "RULES": spaces.Dict( - { - i - + 1: spaces.Dict( - { - "position": i, - "permission": 0, - "source_node_id": 0, - "source_port": 0, - "dest_node_id": 0, - "dest_port": 0, - "protocol": 0, - } - ) - for i in range(self.num_rules) - } - ) + self.default_observation: Dict = { + "RULES": {i+ 1:{ + "position": i, + "permission": 0, + "source_node_id": 0, + "source_port": 0, + "dest_node_id": 0, + "dest_port": 0, + "protocol": 0, + } + for i in range(self.num_rules) } - ) + } def observe(self, state: Dict) -> Dict: - if not self.where: + if self.where is None: return self.default_observation acl_state: Dict = access_from_nested_dict(state, self.where) if acl_state is NOT_PRESENT_IN_STATE: @@ -379,16 +381,16 @@ class AclObservation(AbstractObservation): { "RULE": spaces.Dict( { - i - + 1: spaces.Dict( + i + 1: spaces.Dict( { "position": spaces.Discrete(self.num_rules), "permission": spaces.Discrete(3), - "source_node_id": spaces.Discrete(len(self.nodes) + 1), - "source_port": spaces.Discrete(len(self.ports) + 1), - "dest_node_id": spaces.Discrete(len(self.nodes) + 1), - "dest_port": spaces.Discrete(len(self.ports) + 1), - "protocol": spaces.Discrete(len(self.protocols) + 1), + # adding two to lengths is to account for reserved values 0 (unused) and 1 (any) + "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), + "source_port": spaces.Discrete(len(self.port_to_id) + 2), + "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), + "dest_port": spaces.Discrete(len(self.port_to_id) + 2), + "protocol": spaces.Discrete(len(self.protocol_to_id) + 2), } ) for i in range(self.num_rules) @@ -398,14 +400,96 @@ class AclObservation(AbstractObservation): ) -class ICSObservation(AbstractObservation): - def observe(self, state: Dict) -> Any: - return 0 + + +class NullObservation(AbstractObservation): + def __init__(self, where:Optional[List[str]]=None): + self.default_observation: Dict = {} + + def observe(self, state: Dict) -> Dict: + return {} @property def space(self) -> spaces.Space: - return spaces.Discrete(1) + return spaces.Dict({}) +class ICSObservation(NullObservation): pass + + +class UC2BlueObservation(AbstractObservation): + def __init__( + self, + nodes: List[NodeObservation], + links: List[LinkObservation], + acl: AclObservation, + ics: ICSObservation, + where:Optional[List[str]] = None, + ) -> None: + super().__init__() + self.where: Optional[List[str]] = where + + self.nodes: List[NodeObservation] = nodes + self.links: List[LinkObservation] = links + self.acl: AclObservation = acl + self.ics: ICSObservation = ics + + self.default_observation : Dict = { + "NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)}, + "LINKS": {i+1: l.default_observation for i,l in enumerate(self.links)}, + "ACL": self.acl.default_observation, + "ICS": self.ics.default_observation, + } + + def observe(self, state:Dict) -> Dict: + if self.where is None: + return self.default_observation + + obs = {} + + obs['NODES'] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} + obs['LINKS'] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} + obs['ACL'] = {self.acl.observe(state)} + obs['ICS'] = {self.ics.observe(state)} + + return obs + + @property + def space(self) -> spaces.Space: + return spaces.Dict({ + "NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}), + "LINKS": spaces.Dict({i+1: link.space for i, link in enumerate(self.links)}), + "ACL": self.acl.space, + "ICS": self.ics.space, + }) + + @classmethod + def from_config(cls, config:Dict, sim:Simulation): + nodes = ... + links = ... + acl = ... + ics = ... + new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network']) + return new + + +class UC2RedObservation(AbstractObservation): + def __init__(self, nodes:List[NodeObservation], where:Optional[List[str]] = None) -> None: + super().__init__() + self.where:Optional[List[str]] = where + self.nodes: List[NodeObservation] = nodes + + self.default_observation=...#TODO + + def observe(self, state: Dict) -> Any: + return super().observe(state) + + @property + def space(self) -> spaces.Space: + ... #TODO + + @classmethod + def from_config(cls, config: Dict, sim:Simulation): + ... #TODO class ObservationSpace: """ @@ -422,29 +506,12 @@ class ObservationSpace: # what this class does: # keep a list of observations # create observations for an actor from the config - def __init__( - self, - simulation: Simulation, - nodes: List[NodeObservation] = [], - links: List[LinkObservation] = [], - acl: Optional[AclObservation] = None, - ics: Optional[ICSObservation] = None, - ) -> None: - self.simulation: Simulation = simulation - self.parts: Dict[str, AbstractObservation] = {} + def __init__(self, observation:AbstractObservation) -> None: + self.obs: AbstractObservation = observation - self.nodes: List[NodeObservation] = nodes - self.links: List[LinkObservation] = links - self.acl: Optional[AclObservation] = acl - self.ics: Optional[ICSObservation] = ics - - def observe(self) -> None: - ... + def observe(self, state) -> Dict: + return self.obs.observe(state) @property def space(self) -> None: - ... - - @classmethod - def from_config(self) -> None: - ... + return self.obs.space diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 1db54176..ec778176 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -1,20 +1,17 @@ from abc import ABC, abstractmethod from typing import Any, Dict, List -from pydantic import BaseModel - - -class AbstractReward(BaseModel): - def __call__(self, states: List[Dict]) -> float: - """_summary_ - - :param state: _description_ - :type state: Dict - :return: _description_ - :rtype: float - """ +class AbstractReward(): + def __init__(self): ... + def calculate(self, state:Dict) -> float: + return 0.3 -class RewardFunction(BaseModel): - ... + +class RewardFunction(): + def __init__(self, reward_function:AbstractReward): + self.reward: AbstractReward = reward_function + + def calculate(self, state:Dict) -> float: + return self.reward.calculate(state) diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 47ef4ce9..fcd8b4b3 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -4,3 +4,49 @@ # 3. create actors and configure their actions/observations/rewards/ anything else # 4. Create connection with ARCD GATE # 5. idk + +from primaite.simulator.sim_container import Simulation +from primaite.game.agent.interface import AbstractAgent + +from typing import List + +class PrimaiteSession: + def __init__(self): + self.simulation: Simulation = Simulation() + self.agents:List[AbstractAgent] = [] + self.step_counter:int = 0 + self.episode_counter:int = 0 + + + def step(self): + # currently designed with assumption that all agents act once per step in order + + + for agent in self.agents: + # 3. primaite session asks simulation to provide initial state + # 4. primate session gives state to all agents + # 5. primaite session asks agents to produce an action based on most recent state + sim_state = self.simulation.describe_state() + + # 6. each agent takes most recent state and converts it to CAOS observation + agent_obs = agent.get_obs_from_state(sim_state) + + # 7. meanwhile each agent also takes state and calculates reward + agent_reward = agent.get_reward_from_state(sim_state) + + # 8. each agent takes observation and applies decision rule to observation to create CAOS + # action(such as random, rulebased, or send to GATE) (therefore, converting CAOS action + # to discrete(40) is only necessary for purposes of RL learning, therefore that bit of + # code should live inside of the GATE agent subclass) + # gets action in CAOS format + agent_action = agent.get_action(agent_obs, agent_reward) + # 9. CAOS action is converted into request (extra information might be needed to enrich + # the request, this is what the execution definition is there for) + agent_request = agent.format_request(agent_action) + + # 10. primaite session receives the action from the agents and asks the simulation to apply each + self.simulation.apply_action(agent_request) + + self.simulation.apply_timestep(self.step_counter) + self.step_counter += 1 + From fabd4fd5ddd957ff8abfa3a9d361d04d743fb490 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 4 Oct 2023 09:07:04 +0100 Subject: [PATCH 10/35] Add ACL Action to game layer --- src/primaite/game/agent/actions.py | 48 ++++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index cb7061fc..6c4ae3b2 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -227,11 +227,39 @@ class NodeResetAction(NodeAbstractAction): self.verb = 'reset' class NetworkACLAddRuleAction(AbstractAction): - def __init__(self, manager: "ActionManager", **kwargs) -> None: + def __init__(self, + manager: "ActionManager", + target_router_uuid:str, + max_acl_rules:int, + num_ips:int, + num_ports:int, + num_protocols:int, + **kwargs) -> None: super().__init__(manager=manager) num_permissions = 2 - self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_nics, num_nics, num_ports, num_ports, num_protocols) + self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_ips, num_ips, num_ports, num_ports, num_protocols) + self.target_router_uuid:str = target_router_uuid + def form_request(self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx) -> List[str]: + protocol = self.manager.get_internet_protocol_by_idx(protocol_idx) + src_ip = self.manager.get_ip_address_by_idx(source_ip_idx) + src_port = self.manager.get_port_by_idx(source_port_idx) + dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx) + dst_port = self.manager.get_port_by_idx(dest_port_idx) + return [ + 'network', + 'node', + self.target_router_uuid, + 'acl', + 'add_rule', + permission, + protocol, + src_ip, + src_port, + dst_ip, + dst_port, + position + ] @@ -289,9 +317,14 @@ class ActionManager: max_services_per_node:int = 2, max_nics_per_node:int=8, max_acl_rules:int=10, + protocols:List[str]=['TCP','UDP','ICMP'], + ports:List[str]=['HTTP','DNS','ARP','FTP'], + ip_address_list:Optional[List[str]]=None, act_map:Optional[Dict[int, Dict]]=None) -> None: self.sim: Simulation = sim self.node_uuids:List[str] = node_uuids + self.protocols:List[str] = protocols + self.ports:List[str] = ports action_args = { "num_nodes": len(node_uuids), @@ -299,7 +332,10 @@ class ActionManager: "num_files": max_files_per_folder, "num_services": max_services_per_node, "num_nics": max_nics_per_node, - "num_acl_rules": max_acl_rules} + "num_acl_rules": max_acl_rules, + "num_protocols": len(self.protocols), + "num_ports": len(self.protocols), + "num_ips":} self.actions: Dict[str, AbstractAction] = {} for act_type in actions: self.actions[act_type] = self.__act_class_identifiers[act_type](self, **action_args) @@ -362,8 +398,14 @@ class ActionManager: service_uuids = list(node.services.keys()) return service_uuids[service_idx] if len(service_uuids)>service_idx else None + def get_internet_protocol_by_idx(self, protocol_idx:int) -> str: + # protocol = self.manager.get_internet_protocol_by_idx(protocol_idx) + # src_ip = self.manager.get_ip_address_by_idx(source_ip_idx) + # src_port = self.manager.get_port_by_idx(source_port_idx) + # dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx) + # dst_port = self.manager.get_port_by_idx(dest_port_idx) From 2a8df074b963cefe37f47141af861b6f3ef85238 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 6 Oct 2023 10:36:29 +0100 Subject: [PATCH 11/35] Add network action --- example_config.yaml | 56 ++-- src/primaite/game/agent/actions.py | 385 ++++++++++++++++++--------- src/primaite/game/agent/interface.py | 24 +- src/primaite/game/session.py | 289 +++++++++++++++++++- 4 files changed, 574 insertions(+), 180 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 8cf401cc..b47355c3 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -131,32 +131,36 @@ game_config: action_space: action_list: - - DONOTHING - - NODE_SERVICE_SCAN - - NODE_SERVICE_STOP - # - NODE_SERVICE_START - # - NODE_SERVICE_PAUSE - # - NODE_SERVICE_RESUME - # - NODE_SERVICE_RESTART - # - NODE_SERVICE_DISABLE - # - NODE_SERVICE_ENABLE - # - NODE_FILE_SCAN - # - NODE_FILE_CHECKHASH - # - NODE_FILE_DELETE - # - NODE_FILE_REPAIR - # - NODE_FILE_RESTORE - # - NODE_FOLDER_SCAN - # - NODE_FOLDER_CHECKHASH - # - NODE_FOLDER_REPAIR - # - NODE_FOLDER_RESTORE - # - NODE_OS_SCAN - # - NODE_SHUTDOWN - # - NODE_STARTUP - # - NODE_RESET - # - NETWORK_ACL_ADDRULE - # - NETWORK_ACL_REMOVERULE - # - NETWORK_NIC_ENABLE - - NETWORK_NIC_DISABLE + - type: DONOTHING + - type: NODE_SERVICE_SCAN + - type: NODE_SERVICE_STOP + - type: NODE_SERVICE_START + - type: NODE_SERVICE_PAUSE + - type: NODE_SERVICE_RESUME + - type: NODE_SERVICE_RESTART + - type: NODE_SERVICE_DISABLE + - type: NODE_SERVICE_ENABLE + - type: NODE_FILE_SCAN + - type: NODE_FILE_CHECKHASH + - type: NODE_FILE_DELETE + - type: NODE_FILE_REPAIR + - type: NODE_FILE_RESTORE + - type: NODE_FOLDER_SCAN + - type: NODE_FOLDER_CHECKHASH + - type: NODE_FOLDER_REPAIR + - type: NODE_FOLDER_RESTORE + - type: NODE_OS_SCAN + - type: NODE_SHUTDOWN + - type: NODE_STARTUP + - type: NODE_RESET + - type: NETWORK_ACL_ADDRULE + options: + target_router_ref: router_1 + - type: NETWORK_ACL_REMOVERULE + options: + target_router_ref: router_1 + - type: NETWORK_NIC_ENABLE + - type: NETWORK_NIC_DISABLE action_map: 0: diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 6c4ae3b2..f6f96161 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1,12 +1,13 @@ +import itertools from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple -import itertools - - -from primaite.simulator.sim_container import Simulation from gym import spaces +from primaite.game.session import PrimaiteSession +from primaite.simulator.sim_container import Simulation + + class ExecutionDefiniton(ABC): """ Converter from actions to simulator requests. @@ -23,9 +24,8 @@ class ExecutionDefiniton(ABC): class AbstractAction(ABC): - @abstractmethod - def __init__(self, manager:"ActionManager", **kwargs) -> None: + def __init__(self, manager: "ActionManager", **kwargs) -> None: """ Init method for action. @@ -35,13 +35,12 @@ class AbstractAction(ABC): per node), we need to pass those options to every action that gets created. To pervent verbosity, these parameters are just broadcasted to all actions and the actions can pay attention to the ones that apply. """ - self.name:str = "" + self.name: str = "" """Human-readable action identifier used for printing, logging, and reporting.""" - self.shape = (0,) - """Tuple describing number of options for each parameter of this action. Can be passed to - gym.spaces.MultiDiscrete to form a valid space.""" - self.manager:ActionManager = manager - + self.shape: Dict[str, int] = {} + """Dictionary describing the number of options for each parameter of this action. The keys of this dict must + align with the keyword args of the form_request method.""" + self.manager: ActionManager = manager @abstractmethod def form_request(self) -> List[str]: @@ -50,14 +49,20 @@ class AbstractAction(ABC): class DoNothingAction(AbstractAction): - def __init__(self, manager:"ActionManager", **kwargs) -> None: + def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) self.name = "DONOTHING" - self.shape = (1,) + self.shape: Dict[str, int] = { + "dummy": 1, + } + # This action does not accept any parameters, therefore it technically has a gymnasium shape of Discrete(1), + # i.e. a choice between one option. To make enumerating this action easier, we are adding a 'dummy' paramter + # with one option. This just aids the Action Manager to enumerate all possibilities. def form_request(self) -> List[str]: return ["do_nothing"] + class NodeServiceAbstractAction(AbstractAction): """ Base class for service actions. @@ -65,211 +70,284 @@ class NodeServiceAbstractAction(AbstractAction): Any action which applies to a service and uses node_id and service_id as its only two parameters can inherit from this base class. """ - @abstractmethod - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Tuple[int] = (num_nodes, num_services) - self.verb:str - def form_request(self, node_id:int, service_id:int) -> List[str]: + @abstractmethod + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"node_id": num_nodes, "service_id": num_services} + self.verb: str + + def form_request(self, node_id: int, service_id: int) -> List[str]: node_uuid = self.manager.get_node_uuid_by_idx(node_id) service_uuid = self.manager.get_service_uuid_by_idx(node_id, service_id) if node_uuid is None or service_uuid is None: return ["do_nothing"] - return ['network', 'node', node_uuid, 'services', service_uuid, self.verb] + return ["network", "node", node_uuid, "services", service_uuid, self.verb] + class NodeServiceScanAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "scan" + class NodeServiceStopAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "stop" + class NodeServiceStartAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "start" + class NodeServicePauseAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "pause" + class NodeServiceResumeAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "resume" + class NodeServiceRestartAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "restart" + class NodeServiceDisableAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "disable" + class NodeServiceEnableAction(NodeServiceAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: super().__init__(manager=manager) self.verb = "enable" - class NodeFolderAbstractAction(AbstractAction): @abstractmethod - def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: super().__init__(manager=manager) - self.shape = (num_nodes, num_folders) + self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders} self.verb: str - def form_request(self, node_id:int, folder_id:int) -> List[str]: + def form_request(self, node_id: int, folder_id: int) -> List[str]: node_uuid = self.manager.get_node_uuid_by_idx(node_id) folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) if node_uuid is None or folder_uuid is None: return ["do_nothing"] - return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, self.verb] + return ["network", "node", node_uuid, "file_system", "folder", folder_uuid, self.verb] + class NodeFolderScanAction(NodeFolderAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, **kwargs) - self.verb:str = "scan" + self.verb: str = "scan" + class NodeFolderCheckhashAction(NodeFolderAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, **kwargs) - self.verb:str = "checkhash" + self.verb: str = "checkhash" + class NodeFolderRepairAction(NodeFolderAbstractAction): - def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, **kwargs) - self.verb:str = "repair" + self.verb: str = "repair" + class NodeFolderRestoreAction(NodeFolderAbstractAction): def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, **kwargs) - self.verb:str = "restore" + self.verb: str = "restore" class NodeFileAbstractAction(AbstractAction): @abstractmethod - def __init__(self, manager:"ActionManager", num_nodes:int, num_folders:int, num_files:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager=manager) - self.shape:Tuple[int] = (num_nodes, num_folders, num_files) - self.verb:str + self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders, "file_id": num_files} + self.verb: str - def form_request(self, node_id:int, folder_id:int, file_id:int) -> List[str]: + def form_request(self, node_id: int, folder_id: int, file_id: int) -> List[str]: node_uuid = self.manager.get_node_uuid_by_idx(node_id) folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id) if node_uuid is None or folder_uuid is None or file_uuid is None: return ["do_nothing"] - return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, 'files', file_uuid, self.verb] + return ["network", "node", node_uuid, "file_system", "folder", folder_uuid, "files", file_uuid, self.verb] + class NodeFileScanAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) self.verb = "scan" + class NodeFileCheckhashAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) self.verb = "checkhash" + class NodeFileDeleteAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) self.verb = "delete" + class NodeFileRepairAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) self.verb = "repair" + class NodeFileRestoreAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) self.verb = "restore" + class NodeAbstractAction(AbstractAction): @abstractmethod def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) - self.shape: Tuple[int] = (num_nodes,) + self.shape: Dict[str, int] = {"node_id": num_nodes} self.verb: str - def form_request(self, node_id:int) -> List[str]: + def form_request(self, node_id: int) -> List[str]: node_uuid = self.manager.get_node_uuid_by_idx(node_id) return ["network", "node", node_uuid, self.verb] + class NodeOSScanAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) - self.verb = 'scan' + self.verb = "scan" + class NodeShutdownAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) - self.verb = 'shutdown' + self.verb = "shutdown" + class NodeStartupAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) - self.verb = 'start' + self.verb = "start" + class NodeResetAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) - self.verb = 'reset' + self.verb = "reset" + class NetworkACLAddRuleAction(AbstractAction): - def __init__(self, - manager: "ActionManager", - target_router_uuid:str, - max_acl_rules:int, - num_ips:int, - num_ports:int, - num_protocols:int, - **kwargs) -> None: + def __init__( + self, + manager: "ActionManager", + target_router_uuid: str, + max_acl_rules: int, + num_ips: int, + num_ports: int, + num_protocols: int, + **kwargs, + ) -> None: super().__init__(manager=manager) num_permissions = 2 - self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_ips, num_ips, num_ports, num_ports, num_protocols) - self.target_router_uuid:str = target_router_uuid + self.shape: Dict[str, int] = { + "position": max_acl_rules, + "permission": num_permissions, + "source_ip_idx": num_ips, + "dest_ip_idx": num_ips, + "source_port_idx": num_ports, + "dest_port_idx": num_ports, + "protocol_idx": num_protocols, + } + self.target_router_uuid: str = target_router_uuid - def form_request(self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx) -> List[str]: + def form_request( + self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx + ) -> List[str]: protocol = self.manager.get_internet_protocol_by_idx(protocol_idx) src_ip = self.manager.get_ip_address_by_idx(source_ip_idx) src_port = self.manager.get_port_by_idx(source_port_idx) dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx) dst_port = self.manager.get_port_by_idx(dest_port_idx) return [ - 'network', - 'node', + "network", + "node", self.target_router_uuid, - 'acl', - 'add_rule', + "acl", + "add_rule", permission, protocol, src_ip, src_port, dst_ip, dst_port, - position + position, ] +class NetworkACLRemoveRuleAction(AbstractAction): + def __init__(self, manager: "ActionManager", target_router_uuid: str, max_acl_rules: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"position": max_acl_rules} + self.target_router_uuid: str = target_router_uuid + def form_request(self, position: int) -> List[str]: + return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position] + + +class NetworkNICEnableAction(AbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} + + def form_request(self, node_id: int, nic_id: int) -> List[str]: + return [ + "network", + "node", + self.manager.get_node_uuid_by_idx(node_idx=node_id), + "nic", + self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), + "enable", + ] + + +class NetworkNICDisableAction(AbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} + + def form_request(self, node_id: int, nic_id: int) -> List[str]: + return [ + "network", + "node", + self.manager.get_node_uuid_by_idx(node_idx=node_id), + "nic", + self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), + "disable", + ] class ActionManager: # let the action manager handle the conversion of action spaces into a single discrete integer space. # - # when action space is created, it will take subspaces and generate an action map by enumerating all possibilities, # BUT, the action map can be provided in the config, in which case it will use that. @@ -278,69 +356,83 @@ class ActionManager: # 0: DONOTHING # 1: NODE, FILE, SCAN, NODEID=2, FOLDERID=1, FILEID=0 # 2: ...... - __act_class_identifiers:Dict[str,type] = { + __act_class_identifiers: Dict[str, type] = { "DONOTHING": DoNothingAction, "NODE_SERVICE_SCAN": NodeServiceScanAction, "NODE_SERVICE_STOP": NodeServiceStopAction, - # "NODE_SERVICE_START": NodeServiceStartAction, - # "NODE_SERVICE_PAUSE": NodeServicePauseAction, - # "NODE_SERVICE_RESUME": NodeServiceResumeAction, - # "NODE_SERVICE_RESTART": NodeServiceRestartAction, - # "NODE_SERVICE_DISABLE": NodeServiceDisableAction, - # "NODE_SERVICE_ENABLE": NodeServiceEnableAction, - # "NODE_FILE_SCAN": NodeFileScanAction, - # "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, - # "NODE_FILE_DELETE": NodeFileDeleteAction, - # "NODE_FILE_REPAIR": NodeFileRepairAction, - # "NODE_FILE_RESTORE": NodeFileRestoreAction, + "NODE_SERVICE_START": NodeServiceStartAction, + "NODE_SERVICE_PAUSE": NodeServicePauseAction, + "NODE_SERVICE_RESUME": NodeServiceResumeAction, + "NODE_SERVICE_RESTART": NodeServiceRestartAction, + "NODE_SERVICE_DISABLE": NodeServiceDisableAction, + "NODE_SERVICE_ENABLE": NodeServiceEnableAction, + "NODE_FILE_SCAN": NodeFileScanAction, + "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, + "NODE_FILE_DELETE": NodeFileDeleteAction, + "NODE_FILE_REPAIR": NodeFileRepairAction, + "NODE_FILE_RESTORE": NodeFileRestoreAction, "NODE_FOLDER_SCAN": NodeFolderScanAction, - # "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, - # "NODE_FOLDER_REPAIR": NodeFolderRepairAction, - # "NODE_FOLDER_RESTORE": NodeFolderRestoreAction, - # "NODE_OS_SCAN": NodeOSScanAction, - # "NODE_SHUTDOWN": NodeShutdownAction, - # "NODE_STARTUP": NodeStartupAction, - # "NODE_RESET": NodeResetAction, - # "NETWORK_ACL_ADDRULE": NetworkACLAddRuleAction, - # "NETWORK_ACL_REMOVERULE": NetworkACLRemoveRuleAction, - # "NETWORK_NIC_ENABLE": NetworkNICEnable, - # "NETWORK_NIC_DISABLE": NetworkNICDisable, + "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, + "NODE_FOLDER_REPAIR": NodeFolderRepairAction, + "NODE_FOLDER_RESTORE": NodeFolderRestoreAction, + "NODE_OS_SCAN": NodeOSScanAction, + "NODE_SHUTDOWN": NodeShutdownAction, + "NODE_STARTUP": NodeStartupAction, + "NODE_RESET": NodeResetAction, + "NETWORK_ACL_ADDRULE": NetworkACLAddRuleAction, + "NETWORK_ACL_REMOVERULE": NetworkACLRemoveRuleAction, + "NETWORK_NIC_ENABLE": NetworkNICEnableAction, + "NETWORK_NIC_DISABLE": NetworkNICDisableAction, } + def __init__( + self, + session: PrimaiteSession, # reference to session for looking up stuff + actions: List[str], # stores list of actions available to agent + node_uuids: List[str], # allows mapping index to node + max_folders_per_node: int = 2, # allows calculating shape + max_files_per_folder: int = 2, # allows calculating shape + max_services_per_node: int = 2, # allows calculating shape + max_nics_per_node: int = 8, # allows calculating shape + max_acl_rules: int = 10, # allows calculating shape + protocols: List[str] = ["TCP", "UDP", "ICMP"], # allow mapping index to protocol + ports: List[str] = ["HTTP", "DNS", "ARP", "FTP"], # allow mapping index to port + ip_address_list: Optional[List[str]] = None, # to allow us to map an index to an ip address. + act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions + ) -> None: + self.session: PrimaiteSession = session + self.sim: Simulation = self.session.simulation + self.node_uuids: List[str] = node_uuids + self.protocols: List[str] = protocols + self.ports: List[str] = ports - def __init__(self, - sim:Simulation, - actions:List[str], - node_uuids:List[str], - max_folders_per_node:int = 2, - max_files_per_folder:int = 2, - max_services_per_node:int = 2, - max_nics_per_node:int=8, - max_acl_rules:int=10, - protocols:List[str]=['TCP','UDP','ICMP'], - ports:List[str]=['HTTP','DNS','ARP','FTP'], - ip_address_list:Optional[List[str]]=None, - act_map:Optional[Dict[int, Dict]]=None) -> None: - self.sim: Simulation = sim - self.node_uuids:List[str] = node_uuids - self.protocols:List[str] = protocols - self.ports:List[str] = ports + self.ip_address_list: List[str] + if ip_address_list is not None: + self.ip_address_list = ip_address_list + else: + self.ip_address_list = [] + for node_uuid in self.node_uuids: + node_obj = self.sim.network.nodes[node_uuid] + nics = node_obj.nics + for nic_uuid, nic_obj in nics.items(): + self.ip_address_list.append(nic_obj.ip_address) action_args = { "num_nodes": len(node_uuids), - "num_folders":max_folders_per_node, + "num_folders": max_folders_per_node, "num_files": max_files_per_folder, "num_services": max_services_per_node, "num_nics": max_nics_per_node, "num_acl_rules": max_acl_rules, "num_protocols": len(self.protocols), "num_ports": len(self.protocols), - "num_ips":} + "num_ips": len(self.ip_address_list), + } self.actions: Dict[str, AbstractAction] = {} for act_type in actions: self.actions[act_type] = self.__act_class_identifiers[act_type](self, **action_args) - self.action_map:Dict[int, Tuple[str, Dict]] = {} + self.action_map: Dict[int, Tuple[str, Dict]] = {} """ Action mapping that converts an integer to a specific action and parameter choice. @@ -350,21 +442,30 @@ class ActionManager: if act_map is None: self.action_map = self._enumerate_actions() else: - self.action_map = {i:(a['action'], a['options']) for i,a in act_map.items()} + self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()} # make sure all numbers between 0 and N are represented as dict keys in action map assert all([i in self.action_map.keys() for i in range(len(self.action_map))]) - def _enumerate_actions(self,) -> Dict[int, Tuple[AbstractAction, Dict]]: - ... + def _enumerate_actions( + self, + ) -> Dict[int, Tuple[AbstractAction, Dict]]: + all_action_possibilities = [] + for action in self.actions.values(): + param_names = (list(action.shape.keys()),) + num_possibilities = list(action.shape.values()) + possibilities = [range(n) for n in num_possibilities] - def get_action(self, action: int) -> Tuple[str,Dict]: + itertools.product(action.shape.values()) + all_action_possibilities.append((action, {})) + + def get_action(self, action: int) -> Tuple[str, Dict]: """Produce action in CAOS format""" """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" """The caos format is basically a action identifier, followed by parameters stored in a dictionary""" act_identifier, act_options = self.action_map[action] return act_identifier, act_options - def form_request(self, action_identifier:str, action_options:Dict): + def form_request(self, action_identifier: str, action_options: Dict): """Take action in CAOS format and use the execution definition to change it into PrimAITE request format""" act_obj = self.actions[action_identifier] return act_obj.form_request(**action_options) @@ -380,37 +481,57 @@ class ActionManager: node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] folder_uuids = list(node.file_system.folders.keys()) - return folder_uuids[folder_idx] if len(folder_uuids)>folder_idx else None + return folder_uuids[folder_idx] if len(folder_uuids) > folder_idx else None def get_file_uuid_by_idx(self, node_idx, folder_idx, file_idx) -> Optional[str]: node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] folder_uuids = list(node.file_system.folders.keys()) - if len(folder_uuids)<=folder_idx: + if len(folder_uuids) <= folder_idx: return None folder = node.file_system.folders[folder_uuids[folder_idx]] file_uuids = list(folder.files.keys()) - return file_uuids[file_idx] if len(file_uuids)>file_idx else None + return file_uuids[file_idx] if len(file_uuids) > file_idx else None def get_service_uuid_by_idx(self, node_idx, service_idx) -> Optional[str]: node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] service_uuids = list(node.services.keys()) - return service_uuids[service_idx] if len(service_uuids)>service_idx else None + return service_uuids[service_idx] if len(service_uuids) > service_idx else None - def get_internet_protocol_by_idx(self, protocol_idx:int) -> str: + def get_internet_protocol_by_idx(self, protocol_idx: int) -> str: + return self.protocols[protocol_idx] + def get_ip_address_by_idx(self, ip_idx: int) -> str: + return self.ip_address_list[ip_idx] - # protocol = self.manager.get_internet_protocol_by_idx(protocol_idx) - # src_ip = self.manager.get_ip_address_by_idx(source_ip_idx) - # src_port = self.manager.get_port_by_idx(source_port_idx) - # dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx) - # dst_port = self.manager.get_port_by_idx(dest_port_idx) + def get_port_by_idx(self, port_idx: int) -> str: + return self.ports[port_idx] + def get_nic_uuid_by_idx(self, node_idx: int, nic_idx: int) -> str: + node_uuid = self.get_node_uuid_by_idx(node_idx) + node_obj = self.sim.network.nodes[node_uuid] + nics = list(node_obj.nics.keys()) + if len(nics) <= nic_idx: + return None + return nics[nic_idx] + @classmethod + def from_config(cls, session: PrimaiteSession, cfg: Dict) -> "ActionManager": + obj = cls( + session=session, + actions=cfg["action_list"], + node_uuids=cfg["options"]["nodes"], + max_folders_per_node=cfg["options"]["max_folders_per_node"], + max_files_per_folder=cfg["options"]["max_files_per_folder"], + max_services_per_node=cfg["options"]["max_services_per_node"], + max_nics_per_node=cfg["options"]["max_nics_per_node"], + max_acl_rules=cfg["options"]["max_acl_rules"], + max_X=cfg["options"]["max_X"], + protocols=session.options.ports, + ports=session.options.protocols, + ip_address_list=None, + act_map=cfg["action_map"], + ) -class UC2RedActions(AbstractAction): - ... - -class UC2GreenActionSpace(ActionManager): - ... + return obj diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 0e682b60..528c0b1a 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -2,14 +2,16 @@ # That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word # If you disagree, make a comment in the PR review and we can discuss from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Union, TypeAlias +from typing import Any, Dict, List, Optional, TypeAlias, Union + import numpy as np from primaite.game.agent.actions import ActionManager from primaite.game.agent.observations import ObservationSpace from primaite.game.agent.rewards import RewardFunction -ObsType:TypeAlias = Union[Dict, np.ndarray] +ObsType: TypeAlias = Union[Dict, np.ndarray] + class AbstractAgent(ABC): """Base class for scripted and RL agents.""" @@ -28,31 +30,28 @@ class AbstractAgent(ABC): # by for example specifying target ip addresses, or converting a node ID into a uuid self.execution_definition = None - def get_obs_from_state(self, state:Dict) -> ObsType: + def convert_state_to_obs(self, state: Dict) -> ObsType: """ state : dict state directly from simulation.describe_state output : dict state according to CAOS. """ return self.observation_space.observe(state) - def get_reward_from_state(self, state:Dict) -> float: + def calculate_reward_from_state(self, state: Dict) -> float: return self.reward_function.calculate(state) @abstractmethod - def get_action(self, obs:ObsType, reward:float=None): - # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 1-40, + def get_action(self, obs: ObsType, reward: float = None): + # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 0-39, # then use a bespoke conversion to take 1-40 int back into CAOS action - return ('NODE', 'SERVICE', 'SCAN', '', '') + return ("NODE", "SERVICE", "SCAN", "", "") @abstractmethod def format_request(self, action) -> List[str]: # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR """Format action into format expected by the simulator, and apply execution definition if applicable.""" - return ['network', 'nodes', '', 'file_system', 'folder', 'root', 'scan'] - - - + return ["network", "nodes", "", "file_system", "folder", "root", "scan"] class AbstractScriptedAgent(AbstractAgent): @@ -60,10 +59,11 @@ class AbstractScriptedAgent(AbstractAgent): ... + class RandomAgent(AbstractScriptedAgent): """Agent that ignores its observation and acts completely at random.""" - def get_action(self, obs:ObsType, reward:float=None): + def get_action(self, obs: ObsType, reward: float = None): return self.action_space.space.sample() diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index fcd8b4b3..0f88b322 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -5,23 +5,58 @@ # 4. Create connection with ARCD GATE # 5. idk -from primaite.simulator.sim_container import Simulation -from primaite.game.agent.interface import AbstractAgent +from ipaddress import IPv4Address +from typing import Dict, List + +from pydantic import BaseModel + +from primaite.game.agent.actions import ActionManager +from primaite.game.agent.interface import AbstractAgent +from primaite.game.agent.observations import ( + AclObservation, + FileObservation, + FolderObservation, + ICSObservation, + LinkObservation, + NicObservation, + NodeObservation, + NullObservation, + ServiceObservation, + UC2BlueObservation, + UC2RedObservation, +) +from primaite.simulator.network.hardware.base import Link, NIC, Node +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.router import ACLAction, Router +from primaite.simulator.network.hardware.nodes.server import Server +from primaite.simulator.network.hardware.nodes.switch import Switch +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.sim_container import Simulation +from primaite.simulator.system.applications.database_client import DatabaseClient +from primaite.simulator.system.services.database_service import DatabaseService +from primaite.simulator.system.services.dns_client import DNSClient +from primaite.simulator.system.services.dns_server import DNSServer +from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot +from primaite.simulator.system.services.service import Service + + +class PrimaiteSessionOptions(BaseModel): + ports: List[str] + protocols: List[str] -from typing import List class PrimaiteSession: def __init__(self): self.simulation: Simulation = Simulation() - self.agents:List[AbstractAgent] = [] - self.step_counter:int = 0 - self.episode_counter:int = 0 - + self.agents: List[AbstractAgent] = [] + self.step_counter: int = 0 + self.episode_counter: int = 0 + self.options: PrimaiteSessionOptions def step(self): # currently designed with assumption that all agents act once per step in order - for agent in self.agents: # 3. primaite session asks simulation to provide initial state # 4. primate session gives state to all agents @@ -29,10 +64,10 @@ class PrimaiteSession: sim_state = self.simulation.describe_state() # 6. each agent takes most recent state and converts it to CAOS observation - agent_obs = agent.get_obs_from_state(sim_state) + agent_obs = agent.convert_state_to_obs(sim_state) # 7. meanwhile each agent also takes state and calculates reward - agent_reward = agent.get_reward_from_state(sim_state) + agent_reward = agent.calculate_reward_from_state(sim_state) # 8. each agent takes observation and applies decision rule to observation to create CAOS # action(such as random, rulebased, or send to GATE) (therefore, converting CAOS action @@ -50,3 +85,237 @@ class PrimaiteSession: self.simulation.apply_timestep(self.step_counter) self.step_counter += 1 + @classmethod + def from_config(cls, cfg: dict) -> "PrimaiteSession": + sess = cls() + sim = sess.simulation + net = sim.network + + ref_map_nodes: Dict[str, Node] = {} + ref_map_services: Dict[str, Service] = {} + ref_map_links: Dict[str, Link] = {} + + nodes_cfg = cfg["simulation"]["network"]["nodes"] + links_cfg = cfg["simulation"]["network"]["links"] + for node_cfg in nodes_cfg: + node_ref = node_cfg["ref"] + n_type = node_cfg["type"] + if n_type == "computer": + new_node = Computer( + hostname=node_cfg["hostname"], + ip_address=node_cfg["ip_address"], + subnet_mask=node_cfg["subnet_mask"], + default_gateway=node_cfg["default_gateway"], + dns_server=node_cfg["dns_server"], + ) + elif n_type == "server": + new_node = Server( + hostname=node_cfg["hostname"], + ip_address=node_cfg["ip_address"], + subnet_mask=node_cfg["subnet_mask"], + default_gateway=node_cfg["default_gateway"], + dns_server=node_cfg.get("dns_server"), + ) + elif n_type == "switch": + new_node = Switch(hostname=node_cfg["hostname"], num_ports=node_cfg.get("num_ports")) + elif n_type == "router": + new_node = Router(hostname=node_cfg["hostname"], num_ports=node_cfg.get("num_ports")) + if "ports" in node_cfg: + for port_num, port_cfg in node_cfg["ports"].items(): + new_node.configure_port( + port=port_num, ip_address=port_cfg["ip_address"], subnet_mask=port_cfg["subnet_mask"] + ) + if "acl" in node_cfg: + for r_num, r_cfg in node_cfg["acl"].items(): + # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating + # this: 'r_cfg.get('src_port')' + # Port/IPProtocol. TODO Refactor + new_node.acl.add_rule( + action=ACLAction[r_cfg["action"]], + src_port=None if not (p := r_cfg.get("src_port")) else Port[p], + dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p], + protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p], + src_ip_address=r_cfg.get("ip_address"), + dst_ip_address=r_cfg.get("ip_address"), + position=r_num, + ) + else: + print("invalid node type") + if "services" in node_cfg: + for service_cfg in node_cfg["services"]: + service_ref = service_cfg["ref"] + service_type = service_cfg["type"] + service_types_mapping = { + "DNSClient": DNSClient, # key is equal to the 'name' attr of the service class itself. + "DNSServer": DNSServer, + "DatabaseClient": DatabaseClient, + "DatabaseService": DatabaseService, + # 'database_backup': , + "DataManipulationBot": DataManipulationBot, + # 'web_browser' + } + if service_type in service_types_mapping: + new_node.software_manager.install(service_types_mapping[service_type]) + new_service = new_node.software_manager.software[service_type] + ref_map_services[service_ref] = new_service + else: + print(f"service type not found {service_type}") + # service-dependent options + if service_type == "DatabaseClient": + if "options" in service_cfg: + opt = service_cfg["options"] + if "db_server_ip" in opt: + new_service.configure(server_ip_address=IPv4Address(opt["db_server_ip"])) + if service_type == "DNSServer": + if "options" in service_cfg: + opt = service_cfg["options"] + if "domain_mapping" in opt: + for domain, ip in opt["domain_mapping"].items(): + new_service.dns_register(domain, ip) + if "nics" in node_cfg: + for nic_num, nic_cfg in node_cfg["nics"].items(): + new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"])) + + net.add_node(new_node) + new_node.power_on() + ref_map_nodes[node_ref] = new_node.uuid + + # 2. create links between nodes + for link_cfg in links_cfg: + node_a = net.nodes[ref_map_nodes[link_cfg["endpoint_a_ref"]]] + node_b = net.nodes[ref_map_nodes[link_cfg["endpoint_b_ref"]]] + if isinstance(node_a, Switch): + endpoint_a = node_a.switch_ports[link_cfg["endpoint_a_port"]] + else: + endpoint_a = node_a.ethernet_port[link_cfg["endpoint_a_port"]] + if isinstance(node_b, Switch): + endpoint_b = node_b.switch_ports[link_cfg["endpoint_b_port"]] + else: + endpoint_b = node_b.ethernet_port[link_cfg["endpoint_b_port"]] + new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b) + ref_map_links[link_cfg["ref"]] = new_link.uuid + + # 3. create agents + game_cfg = cfg["game_config"] + ports_cfg = game_cfg["ports"] + protocols_cfg = game_cfg["protocols"] + agents_cfg = game_cfg["agents"] + + for agent_cfg in agents_cfg: + agent_ref = agent_cfg["ref"] + agent_type = agent_cfg["type"] + action_space_cfg = agent_cfg["action_space"] + observation_space_cfg = agent_cfg["observation_space"] + reward_function_cfg = agent_cfg["reward_function"] + + # CREATE OBSERVATION SPACE + if observation_space_cfg is None: + obs_space = NullObservation() + elif observation_space_cfg["type"] == "UC2BlueObservation": + node_obs_list = [] + link_obs_list = [] + + # node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses + node_ip_to_index = {} + for node_idx, node_cfg in enumerate(nodes_cfg): + n_ref = node_cfg["ref"] + n_obj = net.nodes[ref_map_nodes[n_ref]] + for nic_uuid, nic_obj in n_obj.nics.items(): + node_ip_to_index[nic_obj.ip_address] = node_idx + 2 + + for node_obs_cfg in observation_space_cfg["options"]["nodes"]: + node_ref = node_obs_cfg["node_ref"] + folder_obs_list = [] + service_obs_list = [] + if "services" in node_obs_cfg: + for service_obs_cfg in node_obs_cfg["services"]: + service_obs_list.append( + ServiceObservation( + where=[ + "network", + "nodes", + ref_map_nodes[node_ref], + "services", + ref_map_services[service_obs_cfg["service_ref"]], + ] + ) + ) + if "folders" in node_obs_cfg: + for folder_obs_cfg in node_obs_cfg["folders"]: + file_obs_list = [] + if "files" in folder_obs_cfg: + for file_obs_cfg in folder_obs_cfg["files"]: + file_obs_list.append( + FileObservation( + where=[ + "network", + "nodes", + ref_map_nodes[node_ref], + "folders", + folder_obs_cfg["folder_name"], + "files", + file_obs_cfg["file_name"], + ] + ) + ) + folder_obs_list.append( + FolderObservation( + where=[ + "network", + "nodes", + ref_map_nodes[node_ref], + "folders", + folder_obs_cfg["folder_name"], + ], + files=file_obs_list, + ) + ) + nic_obs_list = [] + for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys(): + nic_obs_list.append( + NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid]) + ) + node_obs_list.append( + NodeObservation( + where=["network", "nodes", ref_map_nodes[node_ref]], + services=service_obs_list, + folders=folder_obs_list, + nics=nic_obs_list, + logon_status=False, + ) + ) + for link_obs_cfg in observation_space_cfg["options"]["links"]: + link_ref = link_obs_cfg["link_ref"] + link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]])) + + acl_obs = AclObservation( + node_ip_to_id=node_ip_to_index, + ports=game_cfg["ports"], + protocols=game_cfg["ports"], + where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]], + ) + obs_space = UC2BlueObservation( + nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation() + ) + elif observation_space_cfg["type"] == "UC2RedObservation": + obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim) + else: + print("observation space config not specified correctly.") + obs_space = NullObservation() + + # CREATE ACTION SPACE + action_space = ActionManager.from_config(sess, action_space_cfg) + + # CREATE REWARD FUNCTION + + # CREATE AGENT + if agent_type == "GreenWebBrowsingAgent": + ... + elif agent_type == "GATERLAgent": + ... + elif agent_type == "RedDatabaseCorruptingAgent": + ... + else: + print("agent type not found") + + return sess From 3dea9743c355d61b1fd7b7461abe7a79dc77f822 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 6 Oct 2023 20:32:52 +0100 Subject: [PATCH 12/35] Get primaite session step working --- example_config.yaml | 455 +++++++++++++----------- sandbox.ipynb | 195 ++++++---- src/primaite/game/agent/actions.py | 158 +++++--- src/primaite/game/agent/interface.py | 16 +- src/primaite/game/agent/rewards.py | 17 + src/primaite/game/session.py | 34 +- src/primaite/simulator/sim_container.py | 1 + 7 files changed, 520 insertions(+), 356 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index b47355c3..9c75c92e 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -19,20 +19,28 @@ game_config: type: GreenWebBrowsingAgent observation_space: null action_space: - actions: - - type: DONOTHING - nodes: - - node_ref: client_2 - actions: - - type: LOGON - - type: LOGOFF - applications: - # - application_ref: client_2_web_browser - # actions: - # - type: EXECUTE - # execution_definition: - # target_address: arcd.com - reward_function: null + action_list: + - type: DONOTHING + # - type: NODE_LOGON + # - type: NODE_LOGOFF + # - type: NODE_APPLICATION_EXECUTE + # options: + # execution_definition: + # target_address: arcd.com + + options: + nodes: + - node_ref: client_2 + max_folders_per_node: 2 + max_files_per_folder: 2 + max_services_per_node: 2 + max_nics_per_node: 8 + max_acl_rules: 10 + + reward_function: + reward_components: + - type: DUMMY + agent_settings: start_step: 5 frequency: 4 @@ -41,6 +49,7 @@ game_config: - ref: client_1_data_manipulation_red_bot team: RED type: RedDatabaseCorruptingAgent + observation_space: type: UC2RedObservation options: @@ -55,27 +64,56 @@ game_config: - operating_status - health_status folders: {} + action_space: - actions: - - type: DO_NOTHING - network: + action_list: + - type: DONOTHING + # - type: NODE_APPLICATION_EXECUTE + # options: + # execution_definition: + # server_ip: 192.168.1.14 + # payload: "DROP TABLE IF EXISTS user;" + # success_rate: 80% + - type: NODE_FILE_DELETE + - type: NODE_FILE_CORRUPT + # - type: NODE_FOLDER_DELETE + # - type: NODE_FOLDER_CORRUPT + - type: NODE_OS_SCAN + # - type: NODE_LOGON + # - type: NODE_LOGOFF + options: nodes: - node_ref: client_1 - actions: - - type: SCAN - - type: LOGON - - type: LOGOFF - services: - - service_ref: data_manipulation_bot - actions: - - type: COMPROMISE - execution_definition: - server_ip: 192.168.1.14 - payload: "DROP TABLE IF EXISTS user;" - success_rate: 80% - folders: - files: {} - reward_function: null + max_folders_per_node: 2 + max_files_per_folder: 2 + max_services_per_node: 2 + # max_nics_per_node: 8 + # max_acl_rules: 10 + + # actions: + # - type: DO_NOTHING + # network: + # nodes: + # - node_ref: client_1 + # actions: + # - type: SCAN + # - type: LOGON + # - type: LOGOFF + # services: + # - service_ref: data_manipulation_bot + # actions: + # - type: COMPROMISE + # execution_definition: + # server_ip: 192.168.1.14 + # payload: "DROP TABLE IF EXISTS user;" + # success_rate: 80% + # folders: + # files: {} + + reward_function: + reward_components: + - type: DUMMY + agent_settings: # options specific to this particular agent type, basically args of __init__(self) start_step: 25 frequency: 20 @@ -85,8 +123,9 @@ game_config: - ref: defender - team: blue + team: BLUE type: GATERLAgent + observation_space: type: UC2BlueObservation options: @@ -128,7 +167,6 @@ game_config: router_node_ref: router_1 ics: null - action_space: action_list: - type: DONOTHING @@ -164,227 +202,227 @@ game_config: action_map: 0: - - action: DONOTHING + action: DONOTHING options: {} # scan webapp service 1: - - action: NODE_SERVICE_SCAN + action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 # stop webapp service 2: - - action: NODE_SERVICE_STOP + action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 # start webapp service 3: - - action: "NODE_SERVICE_START" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_START" + options: + - node_id: 2 + - service_id: 1 4: - - action: "NODE_SERVICE_PAUSE" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_PAUSE" + options: + - node_id: 2 + - service_id: 1 5: - - action: "NODE_SERVICE_RESUME" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_RESUME" + options: + - node_id: 2 + - service_id: 1 6: - - action: "NODE_SERVICE_RESTART" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_RESTART" + options: + - node_id: 2 + - service_id: 1 7: - - action: "NODE_SERVICE_DISABLE" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_DISABLE" + options: + - node_id: 2 + - service_id: 1 8: - - action: "NODE_SERVICE_ENABLE" - options: - - node_id: 2 - - service_id: 1 + action: "NODE_SERVICE_ENABLE" + options: + - node_id: 2 + - service_id: 1 9: - - action: "NODE_FILE_SCAN" - options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + action: "NODE_FILE_SCAN" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 10: - - action: "NODE_FILE_CHECKHASH" - options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + action: "NODE_FILE_CHECKHASH" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 11: - - action: "NODE_FILE_DELETE" - options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + action: "NODE_FILE_DELETE" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 12: - - action: "NODE_FILE_REPAIR" - options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + action: "NODE_FILE_REPAIR" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 13: - - action: "NODE_FILE_RESTORE" - options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + action: "NODE_FILE_RESTORE" + options: + - node_id: 3 + - folder_id: 1 + - file_id: 1 14: - - action: "NODE_FOLDER_SCAN" - options: - - node_id: 3 - - folder_id: 1 + action: "NODE_FOLDER_SCAN" + options: + - node_id: 3 + - folder_id: 1 15: - - action: "NODE_FOLDER_CHECKHASH" - options: - - node_id: 3 - - folder_id: 1 + action: "NODE_FOLDER_CHECKHASH" + options: + - node_id: 3 + - folder_id: 1 16: - - action: "NODE_FOLDER_REPAIR" - options: - - node_id: 3 - - folder_id: 1 + action: "NODE_FOLDER_REPAIR" + options: + - node_id: 3 + - folder_id: 1 17: - - action: "NODE_FOLDER_RESTORE" - options: - - node_id: 3 - - folder_id: 1 + action: "NODE_FOLDER_RESTORE" + options: + - node_id: 3 + - folder_id: 1 18: - - action: "NODE_OS_SCAN" - options: - - node_id: 3 + action: "NODE_OS_SCAN" + options: + - node_id: 3 19: - - action: "NODE_SHUTDOWN" - options: - - node_id: 6 + action: "NODE_SHUTDOWN" + options: + - node_id: 6 20: - - action: "NODE_STARTUP" - options: - - node_id: 6 + action: "NODE_STARTUP" + options: + - node_id: 6 21: - - action: "NODE_RESET" - options: - - node_id: 6 + action: "NODE_RESET" + options: + - node_id: 6 22: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 6 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 6 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 23: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 5 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 5 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 24: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 4 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 4 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 25: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 3 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 3 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 26: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 2 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 2 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 27: - - action: "NETWORK_ACL_ADDRULE" - options: - - position: 1 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + action: "NETWORK_ACL_ADDRULE" + options: + - position: 1 + - permission: 2 + - source_node_id: ... + - dest_node_id: ... + - source_port_id: ... + - dest_port_id: ... + - protocol_id: ... 28: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 0 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 0 29: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 1 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 1 30: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 2 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 2 31: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 3 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 3 32: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 4 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 4 33: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 5 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 5 34: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 6 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 6 35: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 7 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 7 36: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 8 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 8 37: - - action: "NETWORK_ACL_REMOVERULE" - options: - - position: 9 + action: "NETWORK_ACL_REMOVERULE" + options: + - position: 9 38: - - action: "NETWORK_NIC_DISABLE" - options: - - node_id: 6 - - nic_index: 1 + action: "NETWORK_NIC_DISABLE" + options: + - node_id: 6 + - nic_index: 1 39: - - action: "NETWORK_NIC_ENABLE" - options: - - node_id: 6 - - nic_index: 1 + action: "NETWORK_NIC_ENABLE" + options: + - node_id: 6 + - nic_index: 1 options: nodes: @@ -404,9 +442,10 @@ game_config: max_nics_per_node: 8 max_acl_rules: 10 - reward_function: - # ... + reward_components: + - type: DUMMY + agent_settings: # ... diff --git a/sandbox.ipynb b/sandbox.ipynb index 3ff72170..51849298 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -2,18 +2,9 @@ "cells": [ { "cell_type": "code", - "execution_count": 13, + "execution_count": 1, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "The autoreload extension is already loaded. To reload it, use:\n", - " %reload_ext autoreload\n" - ] - } - ], + "outputs": [], "source": [ "%load_ext autoreload\n", "%autoreload 2" @@ -21,28 +12,102 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.game.session import PrimaiteSession\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import itertools" + ] + }, + { + "cell_type": "code", + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from primaite.game.session import PrimaiteSession\n", "from primaite.simulator.sim_container import Simulation\n", "from primaite.game.agent.interface import AbstractAgent\n", - "from primaite.simulator.network.networks import arcd_uc2_network\n" + "from primaite.simulator.network.networks import arcd_uc2_network\n", + "import yaml\n" ] }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ - "sess = PrimaiteSession()" + "with open('example_config.yaml', 'r') as file:\n", + " cfg = yaml.safe_load(file)" ] }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-06 19:05:49,548: Added node 387fba92-e5ff-4ead-b525-1872091935ad to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,557: Added node a808ea99-5c8b-42c4-8e38-bf406ceb1f87 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,562: Added node 922c77bb-096a-4236-9e1e-a44da15c1718 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,579: Added node f11cc63b-537c-4813-be09-a1f5597dfe14 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,591: Added node a866b811-efa2-41cc-adc0-a4752f40a0b8 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,607: Added node a01c22b8-cdfb-4105-a8d0-c67c53b3d08b to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,635: Added node 217074fc-021e-4b19-94db-3bc2d5f15d49 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,641: Added node 28db0167-0621-4fdb-9e2b-65e25a91a101 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,648: Added node e754e649-7ba3-4f80-8621-906255cf8749 to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n", + "2023-10-06 19:05:49,657: Added node 65508b30-defa-46c8-af44-9f0c4c0ae59d to Network 35300ca7-ca53-41a4-b617-1f64c7645e52\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "service type not found DatabaseBackup\n", + "service type not found WebBrowser\n" + ] + } + ], + "source": [ + "sess = PrimaiteSession.from_config(cfg)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[]" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sess.agents" + ] + }, + { + "cell_type": "code", + "execution_count": 8, "metadata": {}, "outputs": [], "source": [ @@ -51,7 +116,16 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "sess.step()" + ] + }, + { + "cell_type": "code", + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -74,27 +148,9 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-02 15:10:20,422: Added node 6abb7664-4d17-45ff-a3c7-dbcccffcfd6d to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,424: Added node 3edbc521-3c80-47e3-8017-dbc38fb00a73 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,428: Added node 94457fb9-04a1-4dc1-9ff7-b64df0da7424 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,432: Added node 0d311d72-139c-41bf-aef7-fa9b01b124d7 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,439: Added node 6161e785-f377-48de-aa4f-20d3646da635 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,444: Added node 55a9e9f8-ee3a-4c28-9b6d-c0c0d78a3f6a to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,447: Added node 2f04ca45-3439-489a-81f7-41cca5ae8adc to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,531: Added node 98660c30-8e48-4b96-967a-d62ca71b4d6d to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,545: Added node 1a184184-b204-40de-986a-4d7459036dbe to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,551: Added node 17b92b9a-6805-4677-85f7-c0c0521a6e25 to Network 045a3114-4aac-4687-a10e-432cfd138325\n", - "2023-10-02 15:10:20,555::ERROR::primaite.simulator.network.hardware.base::175::NIC da:f3:1b:87:24:20/192.168.10.110 cannot be enabled as it is not connected to a Link\n" - ] - } - ], + "outputs": [], "source": [ "router_1 = Router(hostname=\"router_1\", num_ports=5)\n", "router_1.power_on()\n", @@ -261,7 +317,7 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -270,7 +326,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -279,7 +335,7 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -305,7 +361,7 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -315,7 +371,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -324,26 +380,9 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "['network',\n", - " 'node',\n", - " '6161e785-f377-48de-aa4f-20d3646da635',\n", - " 'file_system',\n", - " 'folder',\n", - " '5aefe92b-923c-4684-b3bf-e78dd18d4771',\n", - " 'scan']" - ] - }, - "execution_count": 24, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "my_trial_act" ] @@ -455,8 +494,8 @@ " session = cls()\n", " with open(cfg_path, 'r') as file:\n", " conf = yaml.safe_load(file)\n", - " \n", - " #1. create nodes \n", + "\n", + " #1. create nodes\n", " sim = Simulation()\n", " net = sim.network\n", " nodes_cfg = conf['simulation']['network']['nodes']\n", @@ -465,15 +504,15 @@ " node_ref = node_cfg['ref']\n", " n_type = node_cfg['type']\n", " if n_type == 'computer':\n", - " new_node = Computer(hostname = node_cfg['hostname'], \n", - " ip_address = node_cfg['ip_address'], \n", - " subnet_mask = node_cfg['subnet_mask'], \n", + " new_node = Computer(hostname = node_cfg['hostname'],\n", + " ip_address = node_cfg['ip_address'],\n", + " subnet_mask = node_cfg['subnet_mask'],\n", " default_gateway = node_cfg['default_gateway'],\n", " dns_server = node_cfg['dns_server'])\n", " elif n_type == 'server':\n", - " new_node = Server(hostname = node_cfg['hostname'], \n", - " ip_address = node_cfg['ip_address'], \n", - " subnet_mask = node_cfg['subnet_mask'], \n", + " new_node = Server(hostname = node_cfg['hostname'],\n", + " ip_address = node_cfg['ip_address'],\n", + " subnet_mask = node_cfg['subnet_mask'],\n", " default_gateway = node_cfg['default_gateway'],\n", " dns_server = node_cfg.get('dns_server'))\n", " elif n_type == 'switch':\n", @@ -484,12 +523,12 @@ " num_ports = node_cfg.get('num_ports'))\n", " if 'ports' in node_cfg:\n", " for port_num, port_cfg in node_cfg['ports'].items():\n", - " new_node.configure_port(port=port_num, \n", + " new_node.configure_port(port=port_num,\n", " ip_address=port_cfg['ip_address'],\n", " subnet_mask=port_cfg['subnet_mask'])\n", " if 'acl' in node_cfg:\n", " for r_num, r_cfg in node_cfg['acl'].items():\n", - " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating \n", + " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating\n", " # this: 'r_cfg.get('src_port')'\n", " # Port/IPProtocol. TODO Refactor\n", " new_node.acl.add_rule(\n", @@ -570,15 +609,15 @@ " action_space_cfg = agent_cfg['action_space']\n", " observation_space_cfg = agent_cfg['observation_space']\n", " reward_function_cfg = agent_cfg['reward_function']\n", - " \n", + "\n", " # CREATE OBSERVATION SPACE\n", " if observation_space_cfg is None:\n", " obs_space = NullObservation()\n", " elif observation_space_cfg['type'] == 'UC2BlueObservation':\n", " node_obs_list = []\n", " link_obs_list = []\n", - " \n", - " \n", + "\n", + "\n", " #node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses\n", " node_ip_to_index = {}\n", " for node_idx, node_cfg in enumerate(nodes_cfg):\n", @@ -587,8 +626,8 @@ " for nic_uuid, nic_obj in n_obj.nics.items():\n", " node_ip_to_index[nic_obj.ip_address] = node_idx + 2\n", "\n", - " \n", - " \n", + "\n", + "\n", " for node_obs_cfg in observation_space_cfg['options']['nodes']:\n", " node_ref = node_obs_cfg['node_ref']\n", " folder_obs_list = []\n", @@ -618,9 +657,9 @@ " else:\n", " print(\"observation space config not specified correctly.\")\n", " obs_space = NullObservation()\n", - " \n", + "\n", " # CREATE ACTION SPACE\n", - " \n", + "\n", "\n", "\n", " # CREATE REWARD FUNCTION\n", diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index f6f96161..3f674fbb 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1,12 +1,14 @@ import itertools from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from gym import spaces -from primaite.game.session import PrimaiteSession from primaite.simulator.sim_container import Simulation +if TYPE_CHECKING: + from primaite.game.session import PrimaiteSession + class ExecutionDefiniton(ABC): """ @@ -59,7 +61,7 @@ class DoNothingAction(AbstractAction): # i.e. a choice between one option. To make enumerating this action easier, we are adding a 'dummy' paramter # with one option. This just aids the Action Manager to enumerate all possibilities. - def form_request(self) -> List[str]: + def form_request(self, **kwargs) -> List[str]: return ["do_nothing"] @@ -86,56 +88,56 @@ class NodeServiceAbstractAction(AbstractAction): class NodeServiceScanAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "scan" class NodeServiceStopAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "stop" class NodeServiceStartAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "start" class NodeServicePauseAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "pause" class NodeServiceResumeAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "resume" class NodeServiceRestartAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "restart" class NodeServiceDisableAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "disable" class NodeServiceEnableAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: - super().__init__(manager=manager) + def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "enable" class NodeFolderAbstractAction(AbstractAction): @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders} self.verb: str @@ -149,26 +151,26 @@ class NodeFolderAbstractAction(AbstractAction): class NodeFolderScanAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, **kwargs) + def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "scan" class NodeFolderCheckhashAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, **kwargs) + def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "checkhash" class NodeFolderRepairAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, **kwargs) + def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "repair" class NodeFolderRestoreAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, **kwargs) + def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "restore" @@ -190,34 +192,40 @@ class NodeFileAbstractAction(AbstractAction): class NodeFileScanAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "scan" class NodeFileCheckhashAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "checkhash" class NodeFileDeleteAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "delete" class NodeFileRepairAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "repair" class NodeFileRestoreAction(NodeFileAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: - super().__init__(manager, num_nodes, num_folders, num_files, **kwargs) + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "restore" +class NodeFileCorruptAction(NodeFileAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) + self.verb = "corrupt" + + class NodeAbstractAction(AbstractAction): @abstractmethod def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: @@ -232,25 +240,25 @@ class NodeAbstractAction(AbstractAction): class NodeOSScanAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) + super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "scan" class NodeShutdownAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) + super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "shutdown" class NodeStartupAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) + super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "start" class NodeResetAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: - super().__init__(manager=manager) + super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "reset" @@ -371,6 +379,7 @@ class ActionManager: "NODE_FILE_DELETE": NodeFileDeleteAction, "NODE_FILE_REPAIR": NodeFileRepairAction, "NODE_FILE_RESTORE": NodeFileRestoreAction, + "NODE_FILE_CORRUPT": NodeFileCorruptAction, "NODE_FOLDER_SCAN": NodeFolderScanAction, "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, "NODE_FOLDER_REPAIR": NodeFolderRepairAction, @@ -387,7 +396,7 @@ class ActionManager: def __init__( self, - session: PrimaiteSession, # reference to session for looking up stuff + session: "PrimaiteSession", # reference to session for looking up stuff actions: List[str], # stores list of actions available to agent node_uuids: List[str], # allows mapping index to node max_folders_per_node: int = 2, # allows calculating shape @@ -400,7 +409,7 @@ class ActionManager: ip_address_list: Optional[List[str]] = None, # to allow us to map an index to an ip address. act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions ) -> None: - self.session: PrimaiteSession = session + self.session: "PrimaiteSession" = session self.sim: Simulation = self.session.simulation self.node_uuids: List[str] = node_uuids self.protocols: List[str] = protocols @@ -417,7 +426,8 @@ class ActionManager: for nic_uuid, nic_obj in nics.items(): self.ip_address_list.append(nic_obj.ip_address) - action_args = { + # action_args are settings which are applied to the action space as a whole. + global_action_args = { "num_nodes": len(node_uuids), "num_folders": max_folders_per_node, "num_files": max_files_per_folder, @@ -427,10 +437,21 @@ class ActionManager: "num_protocols": len(self.protocols), "num_ports": len(self.protocols), "num_ips": len(self.ip_address_list), + "max_acl_rules":max_acl_rules, + "max_nics_per_node": max_nics_per_node, } self.actions: Dict[str, AbstractAction] = {} - for act_type in actions: - self.actions[act_type] = self.__act_class_identifiers[act_type](self, **action_args) + for act_spec in actions: + # each action is provided into the action space config like this: + # - type: ACTION_TYPE + # options: + # option_1: value1 + # option_2: value2 + # where `type` decides which AbstractAction subclass should be used + # and `options` is an optional dict of options to pass to the init method of the action class + act_type = act_spec.get('type') + act_options = act_spec.get('options', {}) + self.actions[act_type] = self.__act_class_identifiers[act_type](self, **global_action_args, **act_options) self.action_map: Dict[int, Tuple[str, Dict]] = {} """ @@ -448,15 +469,41 @@ class ActionManager: def _enumerate_actions( self, - ) -> Dict[int, Tuple[AbstractAction, Dict]]: + ) -> Dict[int, Tuple[str, Dict]]: + """Generate a list of all the possible actions that could be taken. + + This enumerates all actions all combinations of parametes you could choose for those actions. The output + of this function is intended to populate the self.action_map parameter in the situation where the user provides + a list of action types, but doesn't specify any subset of actions that should be made available to the agent. + + The enumeration relies on the Actions' `shape` attribute. + + :return: An action map maps consecutive integers to a combination of Action type and parameter choices. + An example output could be: + {0: ("DONOTHING", {'dummy': 0}), + 1: ("NODE_OS_SCAN", {'node_id': 0}), + 2: ("NODE_OS_SCAN", {'node_id': 1}), + 3: ("NODE_FOLDER_SCAN", {'node_id:0, folder_id:0}), + ... #etc... + } + :rtype: Dict[int, Tuple[AbstractAction, Dict]] + """ all_action_possibilities = [] - for action in self.actions.values(): - param_names = (list(action.shape.keys()),) + for act_name, action in self.actions.items(): + param_names = list(action.shape.keys()) num_possibilities = list(action.shape.values()) possibilities = [range(n) for n in num_possibilities] - itertools.product(action.shape.values()) - all_action_possibilities.append((action, {})) + param_combinations = list(itertools.product(*possibilities)) + all_action_possibilities.extend( + [ + ( + act_name, {param_names[i]:param_combinations[j][i] for i in range(len(param_names))} + ) for j in range(len(param_combinations))] + ) + + return {i:p for i,p in enumerate(all_action_possibilities)} + def get_action(self, action: int) -> Tuple[str, Dict]: """Produce action in CAOS format""" @@ -517,21 +564,16 @@ class ActionManager: return nics[nic_idx] @classmethod - def from_config(cls, session: PrimaiteSession, cfg: Dict) -> "ActionManager": + def from_config(cls, session: "PrimaiteSession", cfg: Dict) -> "ActionManager": obj = cls( session=session, actions=cfg["action_list"], - node_uuids=cfg["options"]["nodes"], - max_folders_per_node=cfg["options"]["max_folders_per_node"], - max_files_per_folder=cfg["options"]["max_files_per_folder"], - max_services_per_node=cfg["options"]["max_services_per_node"], - max_nics_per_node=cfg["options"]["max_nics_per_node"], - max_acl_rules=cfg["options"]["max_acl_rules"], - max_X=cfg["options"]["max_X"], - protocols=session.options.ports, - ports=session.options.protocols, + # node_uuids=cfg["options"]["node_uuids"], + **cfg['options'], + protocols=session.options.protocols, + ports=session.options.ports, ip_address_list=None, - act_map=cfg["action_map"], + act_map=cfg.get("action_map"), ) return obj diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 528c0b1a..4fd52d96 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -2,7 +2,7 @@ # That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word # If you disagree, make a comment in the PR review and we can discuss from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, TypeAlias, Union +from typing import Any, Dict, List, Optional, Tuple, TypeAlias, Union import numpy as np @@ -41,17 +41,17 @@ class AbstractAgent(ABC): return self.reward_function.calculate(state) @abstractmethod - def get_action(self, obs: ObsType, reward: float = None): + def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 0-39, # then use a bespoke conversion to take 1-40 int back into CAOS action - return ("NODE", "SERVICE", "SCAN", "", "") + return ("DO_NOTHING", {} ) - @abstractmethod - def format_request(self, action) -> List[str]: + def format_request(self, action:Tuple[str,Dict], options:Dict[str, int]) -> List[str]: # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR """Format action into format expected by the simulator, and apply execution definition if applicable.""" - return ["network", "nodes", "", "file_system", "folder", "root", "scan"] + request = self.action_space.form_request(action_identifier=action, action_options=options) + return request class AbstractScriptedAgent(AbstractAgent): @@ -63,8 +63,8 @@ class AbstractScriptedAgent(AbstractAgent): class RandomAgent(AbstractScriptedAgent): """Agent that ignores its observation and acts completely at random.""" - def get_action(self, obs: ObsType, reward: float = None): - return self.action_space.space.sample() + def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: + return self.action_space.get_action(self.action_space.space.sample()) class AbstractGATEAgent(AbstractAgent): diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index ec778176..a4ceb2dd 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -5,13 +5,30 @@ class AbstractReward(): def __init__(self): ... + @abstractmethod def calculate(self, state:Dict) -> float: return 0.3 +class DummyReward(AbstractReward): + + def calculate(self, state: Dict) -> float: + return -0.1 class RewardFunction(): + __rew_class_identifiers:Dict[str,type[AbstractReward]] = { + "DUMMY" : DummyReward + } def __init__(self, reward_function:AbstractReward): self.reward: AbstractReward = reward_function def calculate(self, state:Dict) -> float: return self.reward.calculate(state) + + @classmethod + def from_config(cls, cfg:Dict) -> "RewardFunction": + for rew_component_cfg in cfg['reward_components']: + rew_type = rew_component_cfg['type'] + rew_component = cls.__rew_class_identifiers[rew_type]() + new = cls(reward_function=rew_component) + return new + diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 0f88b322..46e834d6 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -11,7 +11,7 @@ from typing import Dict, List from pydantic import BaseModel from primaite.game.agent.actions import ActionManager -from primaite.game.agent.interface import AbstractAgent +from primaite.game.agent.interface import AbstractAgent, RandomAgent from primaite.game.agent.observations import ( AclObservation, FileObservation, @@ -25,6 +25,7 @@ from primaite.game.agent.observations import ( UC2BlueObservation, UC2RedObservation, ) +from primaite.game.agent.rewards import RewardFunction from primaite.simulator.network.hardware.base import Link, NIC, Node from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.router import ACLAction, Router @@ -74,10 +75,10 @@ class PrimaiteSession: # to discrete(40) is only necessary for purposes of RL learning, therefore that bit of # code should live inside of the GATE agent subclass) # gets action in CAOS format - agent_action = agent.get_action(agent_obs, agent_reward) + agent_action, action_options = agent.get_action(agent_obs, agent_reward) # 9. CAOS action is converted into request (extra information might be needed to enrich # the request, this is what the execution definition is there for) - agent_request = agent.format_request(agent_action) + agent_request = agent.format_request(agent_action, action_options) # 10. primaite session receives the action from the agents and asks the simulation to apply each self.simulation.apply_action(agent_request) @@ -88,6 +89,10 @@ class PrimaiteSession: @classmethod def from_config(cls, cfg: dict) -> "PrimaiteSession": sess = cls() + sess.options = PrimaiteSessionOptions( + ports = cfg['game_config']['ports'], + protocols = cfg['game_config']['protocols'], + ) sim = sess.simulation net = sim.network @@ -304,13 +309,33 @@ class PrimaiteSession: obs_space = NullObservation() # CREATE ACTION SPACE + action_space_cfg['options']['node_uuids'] = [] + # if a list of nodes is defined, convert them from node references to node UUIDs + for action_node_option in action_space_cfg.get('options',{}).pop('nodes', {}): + if 'node_ref' in action_node_option: + node_uuid = ref_map_nodes[action_node_option['node_ref']] + action_space_cfg['options']['node_uuids'].append(node_uuid) + # Each action space can potentially have a different list of nodes that it can apply to. Therefore, + # we will pass node_uuids as a part of the action space config. + # However, it's not possible to specify the node uuids directly in the config, as they are generated + # dynamically, so we have to translate node references to uuids before passing this config on. + + if 'action_list' in action_space_cfg: + for action_config in action_space_cfg['action_list']: + if 'options' in action_config: + if 'target_router_ref' in action_config['options']: + _target = action_config['options']['target_router_ref'] + action_config['options']['target_router_uuid'] = ref_map_nodes[_target] + action_space = ActionManager.from_config(sess, action_space_cfg) # CREATE REWARD FUNCTION + rew_function = RewardFunction.from_config(reward_function_cfg) # CREATE AGENT if agent_type == "GreenWebBrowsingAgent": - ... + new_agent = RandomAgent(action_space=action_space, observation_space=obs_space, reward_function=rew_function) + sess.agents.append(new_agent) elif agent_type == "GATERLAgent": ... elif agent_type == "RedDatabaseCorruptingAgent": @@ -318,4 +343,5 @@ class PrimaiteSession: else: print("agent type not found") + return sess diff --git a/src/primaite/simulator/sim_container.py b/src/primaite/simulator/sim_container.py index d647b0bc..1df5fe12 100644 --- a/src/primaite/simulator/sim_container.py +++ b/src/primaite/simulator/sim_container.py @@ -27,6 +27,7 @@ class Simulation(SimComponent): am.add_action("network", Action(func=self.network._action_manager)) # pass through domain actions to the domain object am.add_action("domain", Action(func=self.domain._action_manager)) + am.add_action("do_nothing", Action(func=lambda request, context: ())) return am def describe_state(self) -> Dict: From ccb36f84004134fb361cf7df94806a3a7c099851 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Sun, 8 Oct 2023 17:02:54 +0100 Subject: [PATCH 13/35] Change observations to make loading from config better --- example_config.yaml | 41 ++++----- sandbox.ipynb | 112 +++++++++++++++++++++--- src/primaite/game/agent/interface.py | 2 + src/primaite/game/agent/observations.py | 109 +++++++++++++++++++++-- src/primaite/game/session.py | 32 +++++-- 5 files changed, 246 insertions(+), 50 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 9c75c92e..9f679223 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -17,10 +17,12 @@ game_config: - ref: client_1_green_user team: GREEN type: GreenWebBrowsingAgent - observation_space: null + observation_space: + type: UC2GreenObservation action_space: action_list: - type: DONOTHING + # # - type: NODE_LOGON # - type: NODE_LOGOFF # - type: NODE_APPLICATION_EXECUTE @@ -68,6 +70,7 @@ game_config: action_space: action_list: - type: DONOTHING + #8f:1d:3f:32:1c:d6\n", + "2023-10-08 14:40:34,938: SwitchPort 8f:1d:3f:32:1c:d6 connected to Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6\n", + "2023-10-08 14:40:34,939: Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6 up\n", + "2023-10-08 14:40:34,939: Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6 up\n", + "2023-10-08 14:40:34,940: Added link b8070c26-6ad0-4d7e-aed8-c1bcdcf9b438 to connect be:b1:a2:ce:eb:4c/192.168.1.1 and 8f:1d:3f:32:1c:d6\n", + "2023-10-08 14:40:34,942: NIC dc:48:6c:bd:8b:b1/192.168.1.1 connected to Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a\n", + "2023-10-08 14:40:34,943: SwitchPort b2:14:a5:82:c0:7a connected to Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a\n", + "2023-10-08 14:40:34,945: Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a up\n", + "2023-10-08 14:40:34,946: Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a up\n", + "2023-10-08 14:40:34,946: Added link 102f5506-a939-4af7-8ebb-8e173e18283c to connect dc:48:6c:bd:8b:b1/192.168.1.1 and b2:14:a5:82:c0:7a\n", + "2023-10-08 14:40:34,947: SwitchPort 00:9f:54:21:e2:f2 connected to Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10\n", + "2023-10-08 14:40:34,949: Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10 up\n", + "2023-10-08 14:40:34,950: NIC 68:69:bf:51:6c:c0/192.168.1.10 connected to Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10\n", + "2023-10-08 14:40:34,951: Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10 up\n", + "2023-10-08 14:40:34,952: Added link 6136fd05-7a16-4afd-aebd-cdf6e255689b to connect 00:9f:54:21:e2:f2 and 68:69:bf:51:6c:c0/192.168.1.10\n", + "2023-10-08 14:40:34,952: SwitchPort 48:cc:7b:ac:dd:f9 connected to Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12\n", + "2023-10-08 14:40:34,954: Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12 up\n", + "2023-10-08 14:40:34,954: NIC 64:15:7d:f0:cd:ce/192.168.1.12 connected to Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12\n", + "2023-10-08 14:40:34,955: Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12 up\n", + "2023-10-08 14:40:34,956: Added link 02c6f4e4-3674-4189-a5a1-334fa86921f6 to connect 48:cc:7b:ac:dd:f9 and 64:15:7d:f0:cd:ce/192.168.1.12\n", + "2023-10-08 14:40:34,957: SwitchPort e4:e3:bb:bf:9e:04 connected to Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14\n", + "2023-10-08 14:40:34,958: Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14 up\n", + "2023-10-08 14:40:34,959: NIC 81:cd:6e:b8:3d:6c/192.168.1.14 connected to Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14\n", + "2023-10-08 14:40:34,960: Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14 up\n", + "2023-10-08 14:40:34,961: Added link 57e0f89d-265b-4d27-838b-828ae9800688 to connect e4:e3:bb:bf:9e:04 and 81:cd:6e:b8:3d:6c/192.168.1.14\n", + "2023-10-08 14:40:34,962: SwitchPort 71:5f:fc:32:79:9f connected to Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16\n", + "2023-10-08 14:40:34,964: Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16 up\n", + "2023-10-08 14:40:34,965: NIC 29:fa:41:0b:f5:1b/192.168.1.16 connected to Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16\n", + "2023-10-08 14:40:34,966: Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16 up\n", + "2023-10-08 14:40:34,967: Added link 1f382171-5e0d-4a76-9500-27dc68c3c7ee to connect 71:5f:fc:32:79:9f and 29:fa:41:0b:f5:1b/192.168.1.16\n", + "2023-10-08 14:40:34,968: SwitchPort 66:5d:d0:ba:c1:91 connected to Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110\n", + "2023-10-08 14:40:34,969: Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110 up\n", + "2023-10-08 14:40:34,970: NIC 0d:22:07:53:7a:e1/192.168.1.110 connected to Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110\n", + "2023-10-08 14:40:34,971: Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110 up\n", + "2023-10-08 14:40:34,972: Added link d8ea175e-50c8-4597-99bf-ac9001b30c77 to connect 66:5d:d0:ba:c1:91 and 0d:22:07:53:7a:e1/192.168.1.110\n", + "2023-10-08 14:40:34,972: SwitchPort 22:f5:91:5a:bb:b1 connected to Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21\n", + "2023-10-08 14:40:34,974: Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21 up\n", + "2023-10-08 14:40:34,975: NIC 82:e5:30:d9:0e:85/192.168.10.21 connected to Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21\n", + "2023-10-08 14:40:34,976: Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21 up\n", + "2023-10-08 14:40:34,977: Added link 40ba49b9-e334-45ce-93da-a1459b80e9a2 to connect 22:f5:91:5a:bb:b1 and 82:e5:30:d9:0e:85/192.168.10.21\n", + "2023-10-08 14:40:34,978: SwitchPort 70:77:d0:12:cd:a0 connected to Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22\n", + "2023-10-08 14:40:34,980: Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22 up\n", + "2023-10-08 14:40:34,981: NIC ef:20:20:d8:9a:11/192.168.10.22 connected to Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22\n", + "2023-10-08 14:40:34,982: Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22 up\n", + "2023-10-08 14:40:34,982: Added link c36027fe-052f-4eb6-b6c6-10bf817c7ac9 to connect 70:77:d0:12:cd:a0 and ef:20:20:d8:9a:11/192.168.10.22\n", + "2023-10-08 14:40:34,983: SwitchPort 62:da:0d:de:eb:27 connected to Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110\n", + "2023-10-08 14:40:34,985: Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110 up\n", + "2023-10-08 14:40:34,986: NIC b8:2b:a3:f0:18:b9/192.168.10.110 connected to Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110\n", + "2023-10-08 14:40:34,987: Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110 up\n", + "2023-10-08 14:40:34,988: Added link 9469edcd-6b36-4333-b948-3eeccf24abcb to connect 62:da:0d:de:eb:27 and b8:2b:a3:f0:18:b9/192.168.10.110\n" ] }, { @@ -93,7 +153,9 @@ { "data": { "text/plain": [ - "[]" + "[,\n", + " ,\n", + " ]" ] }, "execution_count": 7, @@ -118,7 +180,33 @@ "cell_type": "code", "execution_count": 9, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-08 14:40:35,046: Stepping primaite session. Step counter: 0\n", + "2023-10-08 14:40:35,047: Sending simulation state to agent client_1_green_user\n", + "2023-10-08 14:40:35,049: Getting agent action\n", + "2023-10-08 14:40:35,050: Formatting agent action DONOTHING\n", + "2023-10-08 14:40:35,051: Sending request to simulation: ['do_nothing']\n", + "2023-10-08 14:40:35,052: Sending simulation state to agent client_1_data_manipulation_red_bot\n" + ] + }, + { + "ename": "AttributeError", + "evalue": "'NoneType' object has no attribute 'observe'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", + "\u001b[1;32m/home/cade/repos/PrimAITE/sandbox.ipynb Cell 10\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> 1\u001b[0m sess\u001b[39m.\u001b[39;49mstep()\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/session.py:75\u001b[0m, in \u001b[0;36mPrimaiteSession.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 72\u001b[0m sim_state \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msimulation\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 74\u001b[0m \u001b[39m# 6. each agent takes most recent state and converts it to CAOS observation\u001b[39;00m\n\u001b[0;32m---> 75\u001b[0m agent_obs \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mconvert_state_to_obs(sim_state)\n\u001b[1;32m 77\u001b[0m \u001b[39m# 7. meanwhile each agent also takes state and calculates reward\u001b[39;00m\n\u001b[1;32m 78\u001b[0m agent_reward \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39mcalculate_reward_from_state(sim_state)\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/interface.py:40\u001b[0m, in \u001b[0;36mAbstractAgent.convert_state_to_obs\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mconvert_state_to_obs\u001b[39m(\u001b[39mself\u001b[39m, state: Dict) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m ObsType:\n\u001b[1;32m 36\u001b[0m \u001b[39m \u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 37\u001b[0m \u001b[39m state : dict state directly from simulation.describe_state\u001b[39;00m\n\u001b[1;32m 38\u001b[0m \u001b[39m output : dict state according to CAOS.\u001b[39;00m\n\u001b[1;32m 39\u001b[0m \u001b[39m \"\"\"\u001b[39;00m\n\u001b[0;32m---> 40\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mobservation_space\u001b[39m.\u001b[39;49mobserve(state)\n", + "\u001b[0;31mAttributeError\u001b[0m: 'NoneType' object has no attribute 'observe'" + ] + } + ], "source": [ "sess.step()" ] diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 4fd52d96..6083db6f 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -18,10 +18,12 @@ class AbstractAgent(ABC): def __init__( self, + agent_name: Optional[str], action_space: Optional[ActionManager], observation_space: Optional[ObservationSpace], reward_function: Optional[RewardFunction], ) -> None: + self.agent_name:str = agent_name or "unnamed_agent" self.action_space: Optional[ActionManager] = action_space self.observation_space: Optional[ObservationSpace] = observation_space self.reward_function: Optional[RewardFunction] = reward_function diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index f919a723..21f623fd 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -1,10 +1,13 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, Hashable, List, Optional +from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING from gym import spaces from pydantic import BaseModel +from primaite.game.session import PrimaiteSession from primaite.simulator.sim_container import Simulation +if TYPE_CHECKING: + from primaite.game.session import PrimaiteSession NOT_PRESENT_IN_STATE = object() """ @@ -53,6 +56,15 @@ class AbstractObservation(ABC): """Subclasses must define the shape that they expect""" ... + @abstractmethod + @classmethod + def from_config(cls, config:Dict, session:"PrimaiteSession"): + """Create this observation space component form a serialised format. + + The `session` parameter is for a the PrimaiteSession object that spawns this component. During deserialisation, + a subclass of this class may need to translate from a 'reference' to a UUID. + """ + class FileObservation(AbstractObservation): def __init__(self, where: Optional[List[str]] = None) -> None: @@ -84,6 +96,10 @@ class FileObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict({"health_status": spaces.Discrete(6)}) + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None): + return cls(where=parent_where+["files", config["file_name"]]) + class ServiceObservation(AbstractObservation): default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} @@ -115,6 +131,11 @@ class ServiceObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)}) + @classmethod + def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]=None): + return cls(where=parent_where+["services",session.ref_map_services[config['service_ref']]]) + + class LinkObservation(AbstractObservation): default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} @@ -154,6 +175,10 @@ class LinkObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})}) + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession"): + return cls(where=['network','links', session.ref_map_links[config['link_ref']]]) + class FolderObservation(AbstractObservation): def __init__(self, where: Optional[List[str]] = None, files: List[FileObservation] = []) -> None: @@ -209,6 +234,15 @@ class FolderObservation(AbstractObservation): } ) + @classmethod + def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]): + where = parent_where + ["folders", config['folder_name']] + + file_configs = config["files"] + files = [FileObservation.from_config(config=f, session=session, parent_where=where) for f in file_configs] + + return cls(where=where,files=files) + class NicObservation(AbstractObservation): default_observation: spaces.Space = {"nic_status": 0} @@ -230,6 +264,10 @@ class NicObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict({"nic_status": spaces.Discrete(3)}) + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]): + return cls(where=parent_where + ["NICs", config["nic_uuid"]]) + class NodeObservation(AbstractObservation): def __init__( @@ -310,6 +348,25 @@ class NodeObservation(AbstractObservation): return spaces.Dict(space_shape) + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None): + node_uuid = session.ref_map_nodes[config['node_ref']] + if parent_where is None: + where = ["network", "nodes", node_uuid] + else: + where = parent_where + ["nodes", node_uuid] + + svc_configs = config.get('services', {}) + services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs] + folder_configs = config.get('folders', {}) + folders = [FolderObservation.from_config(config=c,session=session, parent_where=where) for c in folder_configs] + nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys() + nic_configs = [{'nic_uuid':n for n in nic_uuids }] + nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs] + logon_status = config.get('logon_status',False) + cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status) + return super().from_config(config, session) + class AclObservation(AbstractObservation): @@ -399,6 +456,21 @@ class AclObservation(AbstractObservation): } ) + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation": + node_ip_to_idx = {} + for node_idx, node_cfg in enumerate(config['node_order']): + n_ref = node_cfg["node_ref"] + n_obj = session.simulation.network.nodes[session.ref_map_nodes[n_ref]] + for nic_uuid, nic_obj in n_obj.nics.items(): + node_ip_to_idx[nic_obj.ip_address] = node_idx + 2 + + router_uuid = session.ref_map_nodes[config['router_node_ref']] + return cls( + node_ip_to_id=node_ip_to_idx, + ports=session.options.ports, + protocols=session.options.protocols, + where=["network", "nodes", router_uuid]) @@ -413,6 +485,10 @@ class NullObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict({}) + @classmethod + def from_config(cls, cfg:Dict) -> "NullObservation": + return cls() + class ICSObservation(NullObservation): pass @@ -463,11 +539,18 @@ class UC2BlueObservation(AbstractObservation): }) @classmethod - def from_config(cls, config:Dict, sim:Simulation): - nodes = ... - links = ... - acl = ... - ics = ... + def from_config(cls, config:Dict, sess:"PrimaiteSession"): + node_configs = config["nodes"] + nodes = [NodeObservation.from_config(n) for n in node_configs] + + link_configs = config["links"] + links = [LinkObservation.from_config(l) for l in link_configs] + + acl_config = config["acl"] + acl = AclObservation.from_config(acl_config) + + ics_config = config["ics"] + ics = ICSObservation.from_config(ics_config) new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network']) return new @@ -489,8 +572,11 @@ class UC2RedObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, sim:Simulation): + ... #TODO +class UC2GreenObservation(NullObservation): pass + class ObservationSpace: """ Manage the observations of an Actor. @@ -515,3 +601,14 @@ class ObservationSpace: @property def space(self) -> None: return self.obs.space + + @classmethod + def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace": + if config['type'] == "UC2BlueObservation": + return cls(UC2BlueObservation(config['options'])) + elif config['type'] == "UC2RedObservation": + return cls(UC2RedObservation(config['options'])) + elif config['type'] == "UC2GreenObservation": + return cls(UC2GreenObservation(config["options"])) + else: + raise ValueError("Observation space type invalid") diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 46e834d6..f0ae05c6 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -23,6 +23,7 @@ from primaite.game.agent.observations import ( NullObservation, ServiceObservation, UC2BlueObservation, + UC2GreenObservation, UC2RedObservation, ) from primaite.game.agent.rewards import RewardFunction @@ -41,6 +42,10 @@ from primaite.simulator.system.services.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.service import Service +from primaite import getLogger + +_LOGGER = getLogger(__name__) + class PrimaiteSessionOptions(BaseModel): ports: List[str] @@ -55,13 +60,19 @@ class PrimaiteSession: self.episode_counter: int = 0 self.options: PrimaiteSessionOptions + self.ref_map_nodes: Dict[str, Node] = {} + self.ref_map_services: Dict[str, Service] = {} + self.ref_map_links: Dict[str, Link] = {} + def step(self): + _LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}") # currently designed with assumption that all agents act once per step in order for agent in self.agents: # 3. primaite session asks simulation to provide initial state # 4. primate session gives state to all agents # 5. primaite session asks agents to produce an action based on most recent state + _LOGGER.debug(f"Sending simulation state to agent {agent.agent_name}") sim_state = self.simulation.describe_state() # 6. each agent takes most recent state and converts it to CAOS observation @@ -75,14 +86,18 @@ class PrimaiteSession: # to discrete(40) is only necessary for purposes of RL learning, therefore that bit of # code should live inside of the GATE agent subclass) # gets action in CAOS format + _LOGGER.debug(f"Getting agent action") agent_action, action_options = agent.get_action(agent_obs, agent_reward) # 9. CAOS action is converted into request (extra information might be needed to enrich # the request, this is what the execution definition is there for) + _LOGGER.debug(f"Formatting agent action {agent_action}") # maybe too many debug log statements agent_request = agent.format_request(agent_action, action_options) # 10. primaite session receives the action from the agents and asks the simulation to apply each + _LOGGER.debug(f"Sending request to simulation: {agent_request}") self.simulation.apply_action(agent_request) + _LOGGER.debug(f"Initiating simulation step {self.step_counter}") self.simulation.apply_timestep(self.step_counter) self.step_counter += 1 @@ -96,9 +111,9 @@ class PrimaiteSession: sim = sess.simulation net = sim.network - ref_map_nodes: Dict[str, Node] = {} - ref_map_services: Dict[str, Service] = {} - ref_map_links: Dict[str, Link] = {} + sess.ref_map_nodes: Dict[str, Node] = {} + sess.ref_map_services: Dict[str, Service] = {} + sess.ref_map_links: Dict[str, Link] = {} nodes_cfg = cfg["simulation"]["network"]["nodes"] links_cfg = cfg["simulation"]["network"]["links"] @@ -304,6 +319,8 @@ class PrimaiteSession: ) elif observation_space_cfg["type"] == "UC2RedObservation": obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim) + elif observation_space_cfg["type"] == "UC2GreenObservation": + obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{})) else: print("observation space config not specified correctly.") obs_space = NullObservation() @@ -334,12 +351,15 @@ class PrimaiteSession: # CREATE AGENT if agent_type == "GreenWebBrowsingAgent": - new_agent = RandomAgent(action_space=action_space, observation_space=obs_space, reward_function=rew_function) + # TODO: implement non-random agents and fix this parsing + new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) sess.agents.append(new_agent) elif agent_type == "GATERLAgent": - ... + new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) + sess.agents.append(new_agent) elif agent_type == "RedDatabaseCorruptingAgent": - ... + new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) + sess.agents.append(new_agent) else: print("agent type not found") From 081a3e519a94ada17100bfceac907b8530060002 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Sun, 8 Oct 2023 17:57:45 +0100 Subject: [PATCH 14/35] Fix certain observation bugs --- sandbox.ipynb | 167 ++++++++++--------- src/primaite/game/agent/observations.py | 55 ++++--- src/primaite/game/session.py | 203 ++++++++++++------------ 3 files changed, 228 insertions(+), 197 deletions(-) diff --git a/sandbox.ipynb b/sandbox.ipynb index 0191f0ae..b3b3be0d 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -70,66 +70,66 @@ "name": "stderr", "output_type": "stream", "text": [ - "2023-10-08 14:40:34,910: Added node 162d4e70-4aa3-4663-b5cb-1b2127d002c6 to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,914: Added node 0f316f2c-f9b9-4972-9fbe-c69900b12c28 to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,916: Added node 6ab23b82-7182-48ac-9bbe-8315772bc4ff to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,918: Added node 5ebb4df7-b285-44a6-8579-92932206542d to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,920: Added node c99b357f-1874-4220-a428-d8b87d7383bb to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,924: Added node 5f98cbb1-9045-4f6e-91f7-4eaa78437f11 to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,926: Added node 27288000-bbfa-4ef8-a413-3e42cfe2e5d4 to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,930: Added node cc54ae23-c46f-4adc-a211-5c32a0f307ad to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,933: Added node 91d75e63-31c8-4fac-922c-d008a21b14dc to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,935: Added node 6e29546f-ab79-41f6-8b3c-439c03e27ab4 to Network 0449bb3a-abb6-418d-b041-8742cddc9c55\n", - "2023-10-08 14:40:34,937: NIC be:b1:a2:ce:eb:4c/192.168.1.1 connected to Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6\n", - "2023-10-08 14:40:34,938: SwitchPort 8f:1d:3f:32:1c:d6 connected to Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6\n", - "2023-10-08 14:40:34,939: Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6 up\n", - "2023-10-08 14:40:34,939: Link be:b1:a2:ce:eb:4c/192.168.1.1<-->8f:1d:3f:32:1c:d6 up\n", - "2023-10-08 14:40:34,940: Added link b8070c26-6ad0-4d7e-aed8-c1bcdcf9b438 to connect be:b1:a2:ce:eb:4c/192.168.1.1 and 8f:1d:3f:32:1c:d6\n", - "2023-10-08 14:40:34,942: NIC dc:48:6c:bd:8b:b1/192.168.1.1 connected to Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a\n", - "2023-10-08 14:40:34,943: SwitchPort b2:14:a5:82:c0:7a connected to Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a\n", - "2023-10-08 14:40:34,945: Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a up\n", - "2023-10-08 14:40:34,946: Link dc:48:6c:bd:8b:b1/192.168.1.1<-->b2:14:a5:82:c0:7a up\n", - "2023-10-08 14:40:34,946: Added link 102f5506-a939-4af7-8ebb-8e173e18283c to connect dc:48:6c:bd:8b:b1/192.168.1.1 and b2:14:a5:82:c0:7a\n", - "2023-10-08 14:40:34,947: SwitchPort 00:9f:54:21:e2:f2 connected to Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10\n", - "2023-10-08 14:40:34,949: Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10 up\n", - "2023-10-08 14:40:34,950: NIC 68:69:bf:51:6c:c0/192.168.1.10 connected to Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10\n", - "2023-10-08 14:40:34,951: Link 00:9f:54:21:e2:f2<-->68:69:bf:51:6c:c0/192.168.1.10 up\n", - "2023-10-08 14:40:34,952: Added link 6136fd05-7a16-4afd-aebd-cdf6e255689b to connect 00:9f:54:21:e2:f2 and 68:69:bf:51:6c:c0/192.168.1.10\n", - "2023-10-08 14:40:34,952: SwitchPort 48:cc:7b:ac:dd:f9 connected to Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12\n", - "2023-10-08 14:40:34,954: Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12 up\n", - "2023-10-08 14:40:34,954: NIC 64:15:7d:f0:cd:ce/192.168.1.12 connected to Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12\n", - "2023-10-08 14:40:34,955: Link 48:cc:7b:ac:dd:f9<-->64:15:7d:f0:cd:ce/192.168.1.12 up\n", - "2023-10-08 14:40:34,956: Added link 02c6f4e4-3674-4189-a5a1-334fa86921f6 to connect 48:cc:7b:ac:dd:f9 and 64:15:7d:f0:cd:ce/192.168.1.12\n", - "2023-10-08 14:40:34,957: SwitchPort e4:e3:bb:bf:9e:04 connected to Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14\n", - "2023-10-08 14:40:34,958: Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14 up\n", - "2023-10-08 14:40:34,959: NIC 81:cd:6e:b8:3d:6c/192.168.1.14 connected to Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14\n", - "2023-10-08 14:40:34,960: Link e4:e3:bb:bf:9e:04<-->81:cd:6e:b8:3d:6c/192.168.1.14 up\n", - "2023-10-08 14:40:34,961: Added link 57e0f89d-265b-4d27-838b-828ae9800688 to connect e4:e3:bb:bf:9e:04 and 81:cd:6e:b8:3d:6c/192.168.1.14\n", - "2023-10-08 14:40:34,962: SwitchPort 71:5f:fc:32:79:9f connected to Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16\n", - "2023-10-08 14:40:34,964: Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16 up\n", - "2023-10-08 14:40:34,965: NIC 29:fa:41:0b:f5:1b/192.168.1.16 connected to Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16\n", - "2023-10-08 14:40:34,966: Link 71:5f:fc:32:79:9f<-->29:fa:41:0b:f5:1b/192.168.1.16 up\n", - "2023-10-08 14:40:34,967: Added link 1f382171-5e0d-4a76-9500-27dc68c3c7ee to connect 71:5f:fc:32:79:9f and 29:fa:41:0b:f5:1b/192.168.1.16\n", - "2023-10-08 14:40:34,968: SwitchPort 66:5d:d0:ba:c1:91 connected to Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110\n", - "2023-10-08 14:40:34,969: Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110 up\n", - "2023-10-08 14:40:34,970: NIC 0d:22:07:53:7a:e1/192.168.1.110 connected to Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110\n", - "2023-10-08 14:40:34,971: Link 66:5d:d0:ba:c1:91<-->0d:22:07:53:7a:e1/192.168.1.110 up\n", - "2023-10-08 14:40:34,972: Added link d8ea175e-50c8-4597-99bf-ac9001b30c77 to connect 66:5d:d0:ba:c1:91 and 0d:22:07:53:7a:e1/192.168.1.110\n", - "2023-10-08 14:40:34,972: SwitchPort 22:f5:91:5a:bb:b1 connected to Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21\n", - "2023-10-08 14:40:34,974: Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21 up\n", - "2023-10-08 14:40:34,975: NIC 82:e5:30:d9:0e:85/192.168.10.21 connected to Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21\n", - "2023-10-08 14:40:34,976: Link 22:f5:91:5a:bb:b1<-->82:e5:30:d9:0e:85/192.168.10.21 up\n", - "2023-10-08 14:40:34,977: Added link 40ba49b9-e334-45ce-93da-a1459b80e9a2 to connect 22:f5:91:5a:bb:b1 and 82:e5:30:d9:0e:85/192.168.10.21\n", - "2023-10-08 14:40:34,978: SwitchPort 70:77:d0:12:cd:a0 connected to Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22\n", - "2023-10-08 14:40:34,980: Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22 up\n", - "2023-10-08 14:40:34,981: NIC ef:20:20:d8:9a:11/192.168.10.22 connected to Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22\n", - "2023-10-08 14:40:34,982: Link 70:77:d0:12:cd:a0<-->ef:20:20:d8:9a:11/192.168.10.22 up\n", - "2023-10-08 14:40:34,982: Added link c36027fe-052f-4eb6-b6c6-10bf817c7ac9 to connect 70:77:d0:12:cd:a0 and ef:20:20:d8:9a:11/192.168.10.22\n", - "2023-10-08 14:40:34,983: SwitchPort 62:da:0d:de:eb:27 connected to Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110\n", - "2023-10-08 14:40:34,985: Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110 up\n", - "2023-10-08 14:40:34,986: NIC b8:2b:a3:f0:18:b9/192.168.10.110 connected to Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110\n", - "2023-10-08 14:40:34,987: Link 62:da:0d:de:eb:27<-->b8:2b:a3:f0:18:b9/192.168.10.110 up\n", - "2023-10-08 14:40:34,988: Added link 9469edcd-6b36-4333-b948-3eeccf24abcb to connect 62:da:0d:de:eb:27 and b8:2b:a3:f0:18:b9/192.168.10.110\n" + "2023-10-08 17:56:35,831: Added node af2f9c15-ecb4-4b65-b48f-63f12acddb88 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,836: Added node 47158854-0917-4037-a6a2-33dde56a120f to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,840: Added node cba8ce63-8064-4f80-bcfe-95ca65221dfa to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,846: Added node e01e7c2b-02ac-4e2d-b7bb-8bc3b6ea6509 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,857: Added node bd5d85ba-5980-45c7-8b28-020a2cfeba0f to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,863: Added node 39e0e37c-4d72-4c76-93cb-4f9c29651ef4 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,878: Added node 7d1063f9-b5e5-4753-966e-1b630325b266 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,881: Added node d85b6abb-0f9e-4853-af26-c9b410e1cb94 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,884: Added node 63b18888-98aa-4182-a014-02999d095bd0 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,888: Added node f514cf8a-a3f1-46d6-be00-994364241ef4 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", + "2023-10-08 17:56:35,890: NIC 27:a9:09:ed:30:5a/192.168.1.1 connected to Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20\n", + "2023-10-08 17:56:35,891: SwitchPort cb:6f:24:8c:7a:20 connected to Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20\n", + "2023-10-08 17:56:35,893: Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20 up\n", + "2023-10-08 17:56:35,896: Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20 up\n", + "2023-10-08 17:56:35,897: Added link 41d994cb-2976-4aa2-b306-649cef4deb80 to connect 27:a9:09:ed:30:5a/192.168.1.1 and cb:6f:24:8c:7a:20\n", + "2023-10-08 17:56:35,899: NIC 5c:fa:b1:a4:69:ec/192.168.1.1 connected to Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87\n", + "2023-10-08 17:56:35,900: SwitchPort 68:54:d7:42:04:87 connected to Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87\n", + "2023-10-08 17:56:35,901: Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87 up\n", + "2023-10-08 17:56:35,903: Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87 up\n", + "2023-10-08 17:56:35,904: Added link d582a248-e968-40eb-9d1b-67143d729e0c to connect 5c:fa:b1:a4:69:ec/192.168.1.1 and 68:54:d7:42:04:87\n", + "2023-10-08 17:56:35,905: SwitchPort c6:bd:77:78:4b:5d connected to Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10\n", + "2023-10-08 17:56:35,908: Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10 up\n", + "2023-10-08 17:56:35,909: NIC 1c:d9:92:e8:d6:3b/192.168.1.10 connected to Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10\n", + "2023-10-08 17:56:35,911: Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10 up\n", + "2023-10-08 17:56:35,912: Added link 13315780-1fcc-4c85-b94b-ef8f14c88a8a to connect c6:bd:77:78:4b:5d and 1c:d9:92:e8:d6:3b/192.168.1.10\n", + "2023-10-08 17:56:35,913: SwitchPort cd:46:af:c4:33:65 connected to Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12\n", + "2023-10-08 17:56:35,916: Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12 up\n", + "2023-10-08 17:56:35,917: NIC aa:cf:2f:71:13:5b/192.168.1.12 connected to Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12\n", + "2023-10-08 17:56:35,918: Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12 up\n", + "2023-10-08 17:56:35,919: Added link 6c2a80f7-f36d-4df6-ac84-d354e5d517dd to connect cd:46:af:c4:33:65 and aa:cf:2f:71:13:5b/192.168.1.12\n", + "2023-10-08 17:56:35,920: SwitchPort 2c:d2:67:ef:68:a8 connected to Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14\n", + "2023-10-08 17:56:35,923: Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14 up\n", + "2023-10-08 17:56:35,924: NIC e1:09:5e:98:ee:a2/192.168.1.14 connected to Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14\n", + "2023-10-08 17:56:35,925: Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14 up\n", + "2023-10-08 17:56:35,926: Added link 1cfdd4f2-22be-4e69-8f1f-daef8e18f543 to connect 2c:d2:67:ef:68:a8 and e1:09:5e:98:ee:a2/192.168.1.14\n", + "2023-10-08 17:56:35,927: SwitchPort 9b:13:8c:a0:8c:82 connected to Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16\n", + "2023-10-08 17:56:35,929: Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16 up\n", + "2023-10-08 17:56:35,930: NIC cc:c2:84:03:1c:42/192.168.1.16 connected to Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16\n", + "2023-10-08 17:56:35,932: Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16 up\n", + "2023-10-08 17:56:35,933: Added link 031111e1-3b05-49ce-bd1f-2cdf77b210f4 to connect 9b:13:8c:a0:8c:82 and cc:c2:84:03:1c:42/192.168.1.16\n", + "2023-10-08 17:56:35,934: SwitchPort a1:70:9e:43:1c:07 connected to Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110\n", + "2023-10-08 17:56:35,937: Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110 up\n", + "2023-10-08 17:56:35,938: NIC e7:58:3c:ed:f7:37/192.168.1.110 connected to Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110\n", + "2023-10-08 17:56:35,939: Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110 up\n", + "2023-10-08 17:56:35,941: Added link f15884e7-0df6-4fa5-bb72-406cb2bdff45 to connect a1:70:9e:43:1c:07 and e7:58:3c:ed:f7:37/192.168.1.110\n", + "2023-10-08 17:56:35,943: SwitchPort a5:da:f2:03:21:e3 connected to Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21\n", + "2023-10-08 17:56:35,946: Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21 up\n", + "2023-10-08 17:56:35,947: NIC cf:63:9f:62:fe:df/192.168.10.21 connected to Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21\n", + "2023-10-08 17:56:35,948: Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21 up\n", + "2023-10-08 17:56:35,950: Added link cc6767fa-25de-4daa-bd47-37b49b15a881 to connect a5:da:f2:03:21:e3 and cf:63:9f:62:fe:df/192.168.10.21\n", + "2023-10-08 17:56:35,951: SwitchPort eb:5b:86:14:bd:d1 connected to Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22\n", + "2023-10-08 17:56:35,953: Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22 up\n", + "2023-10-08 17:56:35,954: NIC 3d:73:a6:62:97:3a/192.168.10.22 connected to Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22\n", + "2023-10-08 17:56:35,955: Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22 up\n", + "2023-10-08 17:56:35,958: Added link b29563ed-0636-4188-ba28-52b74a04da27 to connect eb:5b:86:14:bd:d1 and 3d:73:a6:62:97:3a/192.168.10.22\n", + "2023-10-08 17:56:35,959: SwitchPort e3:3a:0b:03:0b:8c connected to Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110\n", + "2023-10-08 17:56:35,961: Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110 up\n", + "2023-10-08 17:56:35,963: NIC d4:ff:37:8d:e4:3d/192.168.10.110 connected to Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110\n", + "2023-10-08 17:56:35,963: Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110 up\n", + "2023-10-08 17:56:35,964: Added link f6fe7757-a8c1-4cdc-a2b2-49d247117903 to connect e3:3a:0b:03:0b:8c and d4:ff:37:8d:e4:3d/192.168.10.110\n" ] }, { @@ -153,9 +153,9 @@ { "data": { "text/plain": [ - "[,\n", - " ,\n", - " ]" + "[,\n", + " ,\n", + " ]" ] }, "execution_count": 7, @@ -185,25 +185,44 @@ "name": "stderr", "output_type": "stream", "text": [ - "2023-10-08 14:40:35,046: Stepping primaite session. Step counter: 0\n", - "2023-10-08 14:40:35,047: Sending simulation state to agent client_1_green_user\n", - "2023-10-08 14:40:35,049: Getting agent action\n", - "2023-10-08 14:40:35,050: Formatting agent action DONOTHING\n", - "2023-10-08 14:40:35,051: Sending request to simulation: ['do_nothing']\n", - "2023-10-08 14:40:35,052: Sending simulation state to agent client_1_data_manipulation_red_bot\n" + "2023-10-08 17:56:36,041: Stepping primaite session. Step counter: 0\n", + "2023-10-08 17:56:36,043: Sending simulation state to agent client_1_green_user\n", + "2023-10-08 17:56:36,045: Getting agent action\n", + "2023-10-08 17:56:36,047: Formatting agent action DONOTHING\n", + "2023-10-08 17:56:36,048: Sending request to simulation: ['do_nothing']\n", + "2023-10-08 17:56:36,050: Sending simulation state to agent client_1_data_manipulation_red_bot\n" ] }, { - "ename": "AttributeError", - "evalue": "'NoneType' object has no attribute 'observe'", + "name": "stdout", + "output_type": "stream", + "text": [ + "[]\n", + "[]\n" + ] + }, + { + "ename": "TypeError", + "evalue": "unhashable type: 'DataManipulationBot'", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", + "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", "\u001b[1;32m/home/cade/repos/PrimAITE/sandbox.ipynb Cell 10\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> 1\u001b[0m sess\u001b[39m.\u001b[39;49mstep()\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/session.py:75\u001b[0m, in \u001b[0;36mPrimaiteSession.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 72\u001b[0m sim_state \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msimulation\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 74\u001b[0m \u001b[39m# 6. each agent takes most recent state and converts it to CAOS observation\u001b[39;00m\n\u001b[0;32m---> 75\u001b[0m agent_obs \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mconvert_state_to_obs(sim_state)\n\u001b[1;32m 77\u001b[0m \u001b[39m# 7. meanwhile each agent also takes state and calculates reward\u001b[39;00m\n\u001b[1;32m 78\u001b[0m agent_reward \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39mcalculate_reward_from_state(sim_state)\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/session.py:80\u001b[0m, in \u001b[0;36mPrimaiteSession.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 77\u001b[0m sim_state \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msimulation\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 79\u001b[0m \u001b[39m# 6. each agent takes most recent state and converts it to CAOS observation\u001b[39;00m\n\u001b[0;32m---> 80\u001b[0m agent_obs \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mconvert_state_to_obs(sim_state)\n\u001b[1;32m 82\u001b[0m \u001b[39m# 7. meanwhile each agent also takes state and calculates reward\u001b[39;00m\n\u001b[1;32m 83\u001b[0m agent_reward \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39mcalculate_reward_from_state(sim_state)\n", "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/interface.py:40\u001b[0m, in \u001b[0;36mAbstractAgent.convert_state_to_obs\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mconvert_state_to_obs\u001b[39m(\u001b[39mself\u001b[39m, state: Dict) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m ObsType:\n\u001b[1;32m 36\u001b[0m \u001b[39m \u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 37\u001b[0m \u001b[39m state : dict state directly from simulation.describe_state\u001b[39;00m\n\u001b[1;32m 38\u001b[0m \u001b[39m output : dict state according to CAOS.\u001b[39;00m\n\u001b[1;32m 39\u001b[0m \u001b[39m \"\"\"\u001b[39;00m\n\u001b[0;32m---> 40\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mobservation_space\u001b[39m.\u001b[39;49mobserve(state)\n", - "\u001b[0;31mAttributeError\u001b[0m: 'NoneType' object has no attribute 'observe'" + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:608\u001b[0m, in \u001b[0;36mObservationSpace.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 607\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mobserve\u001b[39m(\u001b[39mself\u001b[39m, state) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Dict:\n\u001b[0;32m--> 608\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mobs\u001b[39m.\u001b[39;49mobserve(state)\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:571\u001b[0m, in \u001b[0;36mUC2RedObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 568\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[1;32m 570\u001b[0m obs \u001b[39m=\u001b[39m {}\n\u001b[0;32m--> 571\u001b[0m obs[\u001b[39m'\u001b[39m\u001b[39mNODES\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m {i\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m: node\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, node \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes)}\n\u001b[1;32m 572\u001b[0m \u001b[39mreturn\u001b[39;00m obs\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:571\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 568\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[1;32m 570\u001b[0m obs \u001b[39m=\u001b[39m {}\n\u001b[0;32m--> 571\u001b[0m obs[\u001b[39m'\u001b[39m\u001b[39mNODES\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m {i\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m: node\u001b[39m.\u001b[39;49mobserve(state) \u001b[39mfor\u001b[39;00m i, node \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes)}\n\u001b[1;32m 572\u001b[0m \u001b[39mreturn\u001b[39;00m obs\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:328\u001b[0m, in \u001b[0;36mNodeObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 326\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)\n\u001b[1;32m 327\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)\n\u001b[0;32m--> 328\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mSERVICES\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: service\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, service \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)}\n\u001b[1;32m 329\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mFOLDERS\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: folder\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, folder \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)}\n\u001b[1;32m 330\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39moperating_status\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m node_state[\u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m]\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:328\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 326\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)\n\u001b[1;32m 327\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)\n\u001b[0;32m--> 328\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mSERVICES\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: service\u001b[39m.\u001b[39;49mobserve(state) \u001b[39mfor\u001b[39;00m i, service \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)}\n\u001b[1;32m 329\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mFOLDERS\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: folder\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, folder \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)}\n\u001b[1;32m 330\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39moperating_status\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m node_state[\u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m]\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:124\u001b[0m, in \u001b[0;36mServiceObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 121\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mwhere \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 122\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[0;32m--> 124\u001b[0m service_state \u001b[39m=\u001b[39m access_from_nested_dict(state, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mwhere)\n\u001b[1;32m 125\u001b[0m \u001b[39mif\u001b[39;00m service_state \u001b[39mis\u001b[39;00m NOT_PRESENT_IN_STATE:\n\u001b[1;32m 126\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", + " \u001b[0;31m[... skipping similar frames: access_from_nested_dict at line 37 (1 times)]\u001b[0m\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", + "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:35\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 33\u001b[0m \u001b[39mreturn\u001b[39;00m dictionary\n\u001b[1;32m 34\u001b[0m k \u001b[39m=\u001b[39m keys\u001b[39m.\u001b[39mpop(\u001b[39m0\u001b[39m)\n\u001b[0;32m---> 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;49;00m \u001b[39min\u001b[39;49;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[1;32m 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", + "\u001b[0;31mTypeError\u001b[0m: unhashable type: 'DataManipulationBot'" ] } ], diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 21f623fd..c5b931ee 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -3,7 +3,6 @@ from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING from gym import spaces from pydantic import BaseModel -from primaite.game.session import PrimaiteSession from primaite.simulator.sim_container import Simulation if TYPE_CHECKING: @@ -56,8 +55,8 @@ class AbstractObservation(ABC): """Subclasses must define the shape that they expect""" ... - @abstractmethod @classmethod + @abstractmethod def from_config(cls, config:Dict, session:"PrimaiteSession"): """Create this observation space component form a serialised format. @@ -132,7 +131,7 @@ class ServiceObservation(AbstractObservation): return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)}) @classmethod - def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]=None): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]=None): return cls(where=parent_where+["services",session.ref_map_services[config['service_ref']]]) @@ -235,7 +234,7 @@ class FolderObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]): where = parent_where + ["folders", config['folder_name']] file_configs = config["files"] @@ -324,7 +323,6 @@ class NodeObservation(AbstractObservation): return self.default_observation obs = {} - obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)} obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)} obs["operating_status"] = node_state["operating_state"] @@ -349,7 +347,7 @@ class NodeObservation(AbstractObservation): return spaces.Dict(space_shape) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None) -> "NodeObservation": node_uuid = session.ref_map_nodes[config['node_ref']] if parent_where is None: where = ["network", "nodes", node_uuid] @@ -361,12 +359,10 @@ class NodeObservation(AbstractObservation): folder_configs = config.get('folders', {}) folders = [FolderObservation.from_config(config=c,session=session, parent_where=where) for c in folder_configs] nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys() - nic_configs = [{'nic_uuid':n for n in nic_uuids }] + nic_configs = [{'nic_uuid':n for n in nic_uuids }] if nic_uuids else [] nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs] logon_status = config.get('logon_status',False) - cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status) - return super().from_config(config, session) - + return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status) class AclObservation(AbstractObservation): @@ -486,7 +482,7 @@ class NullObservation(AbstractObservation): return spaces.Dict({}) @classmethod - def from_config(cls, cfg:Dict) -> "NullObservation": + def from_config(cls, config:Dict, session:Optional["PrimaiteSession"]=None) -> "NullObservation": return cls() class ICSObservation(NullObservation): pass @@ -539,15 +535,15 @@ class UC2BlueObservation(AbstractObservation): }) @classmethod - def from_config(cls, config:Dict, sess:"PrimaiteSession"): + def from_config(cls, config:Dict, session:"PrimaiteSession"): node_configs = config["nodes"] - nodes = [NodeObservation.from_config(n) for n in node_configs] + nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs] link_configs = config["links"] - links = [LinkObservation.from_config(l) for l in link_configs] + links = [LinkObservation.from_config(config=l, session=session) for l in link_configs] acl_config = config["acl"] - acl = AclObservation.from_config(acl_config) + acl = AclObservation.from_config(config=acl_config, session=session) ics_config = config["ics"] ics = ICSObservation.from_config(ics_config) @@ -561,19 +557,30 @@ class UC2RedObservation(AbstractObservation): self.where:Optional[List[str]] = where self.nodes: List[NodeObservation] = nodes - self.default_observation=...#TODO + self.default_observation : Dict = { + "NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)}, + } - def observe(self, state: Dict) -> Any: - return super().observe(state) + def observe(self, state: Dict) -> Dict: + if self.where is None: + return self.default_observation + + obs = {} + obs['NODES'] = {i+1: node.observe(state) for i, node in enumerate(self.nodes)} + return obs @property def space(self) -> spaces.Space: - ... #TODO + return spaces.Dict({ + "NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}), + }) @classmethod - def from_config(cls, config: Dict, sim:Simulation): + def from_config(cls, config: Dict, session: "PrimaiteSession"): + node_configs = config["nodes"] + nodes = [NodeObservation.from_config(config=cfg, session=session) for cfg in node_configs] + return cls(nodes=nodes, where=["network"]) - ... #TODO class UC2GreenObservation(NullObservation): pass @@ -605,10 +612,10 @@ class ObservationSpace: @classmethod def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace": if config['type'] == "UC2BlueObservation": - return cls(UC2BlueObservation(config['options'])) + return cls(UC2BlueObservation.from_config(config.get('options',{}), session=session)) elif config['type'] == "UC2RedObservation": - return cls(UC2RedObservation(config['options'])) + return cls(UC2RedObservation.from_config(config.get('options',{}), session=session)) elif config['type'] == "UC2GreenObservation": - return cls(UC2GreenObservation(config["options"])) + return cls(UC2GreenObservation.from_config(config.get("options",{}), session=session)) else: raise ValueError("Observation space type invalid") diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index f0ae05c6..4bcf26e4 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -21,6 +21,7 @@ from primaite.game.agent.observations import ( NicObservation, NodeObservation, NullObservation, + ObservationSpace, ServiceObservation, UC2BlueObservation, UC2GreenObservation, @@ -177,7 +178,7 @@ class PrimaiteSession: if service_type in service_types_mapping: new_node.software_manager.install(service_types_mapping[service_type]) new_service = new_node.software_manager.software[service_type] - ref_map_services[service_ref] = new_service + sess.ref_map_services[service_ref] = new_service else: print(f"service type not found {service_type}") # service-dependent options @@ -198,12 +199,12 @@ class PrimaiteSession: net.add_node(new_node) new_node.power_on() - ref_map_nodes[node_ref] = new_node.uuid + sess.ref_map_nodes[node_ref] = new_node.uuid # 2. create links between nodes for link_cfg in links_cfg: - node_a = net.nodes[ref_map_nodes[link_cfg["endpoint_a_ref"]]] - node_b = net.nodes[ref_map_nodes[link_cfg["endpoint_b_ref"]]] + node_a = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_a_ref"]]] + node_b = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_b_ref"]]] if isinstance(node_a, Switch): endpoint_a = node_a.switch_ports[link_cfg["endpoint_a_port"]] else: @@ -213,7 +214,7 @@ class PrimaiteSession: else: endpoint_b = node_b.ethernet_port[link_cfg["endpoint_b_port"]] new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b) - ref_map_links[link_cfg["ref"]] = new_link.uuid + sess.ref_map_links[link_cfg["ref"]] = new_link.uuid # 3. create agents game_cfg = cfg["game_config"] @@ -229,108 +230,112 @@ class PrimaiteSession: reward_function_cfg = agent_cfg["reward_function"] # CREATE OBSERVATION SPACE - if observation_space_cfg is None: - obs_space = NullObservation() - elif observation_space_cfg["type"] == "UC2BlueObservation": - node_obs_list = [] - link_obs_list = [] + obs_space=ObservationSpace.from_config(observation_space_cfg, sess) - # node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses - node_ip_to_index = {} - for node_idx, node_cfg in enumerate(nodes_cfg): - n_ref = node_cfg["ref"] - n_obj = net.nodes[ref_map_nodes[n_ref]] - for nic_uuid, nic_obj in n_obj.nics.items(): - node_ip_to_index[nic_obj.ip_address] = node_idx + 2 + """ + # if observation_space_cfg is None: + # obs_space = NullObservation() + # elif observation_space_cfg["type"] == "UC2BlueObservation": + # node_obs_list = [] + # link_obs_list = [] - for node_obs_cfg in observation_space_cfg["options"]["nodes"]: - node_ref = node_obs_cfg["node_ref"] - folder_obs_list = [] - service_obs_list = [] - if "services" in node_obs_cfg: - for service_obs_cfg in node_obs_cfg["services"]: - service_obs_list.append( - ServiceObservation( - where=[ - "network", - "nodes", - ref_map_nodes[node_ref], - "services", - ref_map_services[service_obs_cfg["service_ref"]], - ] - ) - ) - if "folders" in node_obs_cfg: - for folder_obs_cfg in node_obs_cfg["folders"]: - file_obs_list = [] - if "files" in folder_obs_cfg: - for file_obs_cfg in folder_obs_cfg["files"]: - file_obs_list.append( - FileObservation( - where=[ - "network", - "nodes", - ref_map_nodes[node_ref], - "folders", - folder_obs_cfg["folder_name"], - "files", - file_obs_cfg["file_name"], - ] - ) - ) - folder_obs_list.append( - FolderObservation( - where=[ - "network", - "nodes", - ref_map_nodes[node_ref], - "folders", - folder_obs_cfg["folder_name"], - ], - files=file_obs_list, - ) - ) - nic_obs_list = [] - for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys(): - nic_obs_list.append( - NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid]) - ) - node_obs_list.append( - NodeObservation( - where=["network", "nodes", ref_map_nodes[node_ref]], - services=service_obs_list, - folders=folder_obs_list, - nics=nic_obs_list, - logon_status=False, - ) - ) - for link_obs_cfg in observation_space_cfg["options"]["links"]: - link_ref = link_obs_cfg["link_ref"] - link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]])) + # # node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses + # node_ip_to_index = {} + # for node_idx, node_cfg in enumerate(nodes_cfg): + # n_ref = node_cfg["ref"] + # n_obj = net.nodes[ref_map_nodes[n_ref]] + # for nic_uuid, nic_obj in n_obj.nics.items(): + # node_ip_to_index[nic_obj.ip_address] = node_idx + 2 - acl_obs = AclObservation( - node_ip_to_id=node_ip_to_index, - ports=game_cfg["ports"], - protocols=game_cfg["ports"], - where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]], - ) - obs_space = UC2BlueObservation( - nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation() - ) - elif observation_space_cfg["type"] == "UC2RedObservation": - obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim) - elif observation_space_cfg["type"] == "UC2GreenObservation": - obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{})) - else: - print("observation space config not specified correctly.") - obs_space = NullObservation() + # for node_obs_cfg in observation_space_cfg["options"]["nodes"]: + # node_ref = node_obs_cfg["node_ref"] + # folder_obs_list = [] + # service_obs_list = [] + # if "services" in node_obs_cfg: + # for service_obs_cfg in node_obs_cfg["services"]: + # service_obs_list.append( + # ServiceObservation( + # where=[ + # "network", + # "nodes", + # ref_map_nodes[node_ref], + # "services", + # ref_map_services[service_obs_cfg["service_ref"]], + # ] + # ) + # ) + # if "folders" in node_obs_cfg: + # for folder_obs_cfg in node_obs_cfg["folders"]: + # file_obs_list = [] + # if "files" in folder_obs_cfg: + # for file_obs_cfg in folder_obs_cfg["files"]: + # file_obs_list.append( + # FileObservation( + # where=[ + # "network", + # "nodes", + # ref_map_nodes[node_ref], + # "folders", + # folder_obs_cfg["folder_name"], + # "files", + # file_obs_cfg["file_name"], + # ] + # ) + # ) + # folder_obs_list.append( + # FolderObservation( + # where=[ + # "network", + # "nodes", + # ref_map_nodes[node_ref], + # "folders", + # folder_obs_cfg["folder_name"], + # ], + # files=file_obs_list, + # ) + # ) + # nic_obs_list = [] + # for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys(): + # nic_obs_list.append( + # NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid]) + # ) + # node_obs_list.append( + # NodeObservation( + # where=["network", "nodes", ref_map_nodes[node_ref]], + # services=service_obs_list, + # folders=folder_obs_list, + # nics=nic_obs_list, + # logon_status=False, + # ) + # ) + # for link_obs_cfg in observation_space_cfg["options"]["links"]: + # link_ref = link_obs_cfg["link_ref"] + # link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]])) + + # acl_obs = AclObservation( + # node_ip_to_id=node_ip_to_index, + # ports=game_cfg["ports"], + # protocols=game_cfg["ports"], + # where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]], + # ) + # obs_space = UC2BlueObservation( + # nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation() + # ) + # elif observation_space_cfg["type"] == "UC2RedObservation": + # obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim) + # elif observation_space_cfg["type"] == "UC2GreenObservation": + # obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{})) + # else: + # print("observation space config not specified correctly.") + # obs_space = NullObservation() + """ # CREATE ACTION SPACE action_space_cfg['options']['node_uuids'] = [] # if a list of nodes is defined, convert them from node references to node UUIDs for action_node_option in action_space_cfg.get('options',{}).pop('nodes', {}): if 'node_ref' in action_node_option: - node_uuid = ref_map_nodes[action_node_option['node_ref']] + node_uuid = sess.ref_map_nodes[action_node_option['node_ref']] action_space_cfg['options']['node_uuids'].append(node_uuid) # Each action space can potentially have a different list of nodes that it can apply to. Therefore, # we will pass node_uuids as a part of the action space config. @@ -342,7 +347,7 @@ class PrimaiteSession: if 'options' in action_config: if 'target_router_ref' in action_config['options']: _target = action_config['options']['target_router_ref'] - action_config['options']['target_router_uuid'] = ref_map_nodes[_target] + action_config['options']['target_router_uuid'] = sess.ref_map_nodes[_target] action_space = ActionManager.from_config(sess, action_space_cfg) From f68886d5dfa11f5090a37f35c4049d25dfb23870 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 9 Oct 2023 17:29:50 +0100 Subject: [PATCH 15/35] Fix bugged actions --- example_config.yaml | 202 +++++++++--------- src/primaite/game/agent/actions.py | 118 +++++++--- src/primaite/game/agent/observations.py | 58 ++--- .../network/hardware/nodes/router.py | 10 +- .../network/hardware/nodes/switch.py | 11 +- 5 files changed, 230 insertions(+), 169 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 9f679223..f7faf589 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -63,8 +63,8 @@ game_config: services: - service_ref: data_manipulation_bot observations: - - operating_status - - health_status + operating_status + health_status folders: {} action_space: @@ -197,221 +197,221 @@ game_config: 1: action: NODE_SERVICE_SCAN options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 # stop webapp service 2: action: NODE_SERVICE_STOP options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 # start webapp service 3: action: "NODE_SERVICE_START" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 4: action: "NODE_SERVICE_PAUSE" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 5: action: "NODE_SERVICE_RESUME" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 6: action: "NODE_SERVICE_RESTART" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 7: action: "NODE_SERVICE_DISABLE" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 8: action: "NODE_SERVICE_ENABLE" options: - - node_id: 2 - - service_id: 1 + node_id: 2 + service_id: 1 9: action: "NODE_FILE_SCAN" options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + node_id: 3 + folder_id: 1 + file_id: 1 10: action: "NODE_FILE_CHECKHASH" options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + node_id: 3 + folder_id: 1 + file_id: 1 11: action: "NODE_FILE_DELETE" options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + node_id: 3 + folder_id: 1 + file_id: 1 12: action: "NODE_FILE_REPAIR" options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + node_id: 3 + folder_id: 1 + file_id: 1 13: action: "NODE_FILE_RESTORE" options: - - node_id: 3 - - folder_id: 1 - - file_id: 1 + node_id: 3 + folder_id: 1 + file_id: 1 14: action: "NODE_FOLDER_SCAN" options: - - node_id: 3 - - folder_id: 1 + node_id: 3 + folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - - node_id: 3 - - folder_id: 1 + node_id: 3 + folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - - node_id: 3 - - folder_id: 1 + node_id: 3 + folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - - node_id: 3 - - folder_id: 1 + node_id: 3 + folder_id: 1 18: action: "NODE_OS_SCAN" options: - - node_id: 3 + node_id: 3 19: action: "NODE_SHUTDOWN" options: - - node_id: 6 + node_id: 6 20: action: "NODE_STARTUP" options: - - node_id: 6 + node_id: 6 21: action: "NODE_RESET" options: - - node_id: 6 + node_id: 6 22: action: "NETWORK_ACL_ADDRULE" options: - - position: 6 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 7 + dest_ip_id: 1 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 1 23: action: "NETWORK_ACL_ADDRULE" options: - - position: 5 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 8 + dest_ip_id: 1 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 1 24: action: "NETWORK_ACL_ADDRULE" options: - - position: 4 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 7 + dest_ip_id: 3 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 25: action: "NETWORK_ACL_ADDRULE" options: - - position: 3 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 8 + dest_ip_id: 3 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - - position: 2 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 7 + dest_ip_id: 4 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - - position: 1 - - permission: 2 - - source_node_id: ... - - dest_node_id: ... - - source_port_id: ... - - dest_port_id: ... - - protocol_id: ... + position: 1 + permission: 2 + source_ip_id: 8 + dest_ip_id: 4 + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 28: action: "NETWORK_ACL_REMOVERULE" options: - - position: 0 + position: 0 29: action: "NETWORK_ACL_REMOVERULE" options: - - position: 1 + position: 1 30: action: "NETWORK_ACL_REMOVERULE" options: - - position: 2 + position: 2 31: action: "NETWORK_ACL_REMOVERULE" options: - - position: 3 + position: 3 32: action: "NETWORK_ACL_REMOVERULE" options: - - position: 4 + position: 4 33: action: "NETWORK_ACL_REMOVERULE" options: - - position: 5 + position: 5 34: action: "NETWORK_ACL_REMOVERULE" options: - - position: 6 + position: 6 35: action: "NETWORK_ACL_REMOVERULE" options: - - position: 7 + position: 7 36: action: "NETWORK_ACL_REMOVERULE" options: - - position: 8 + position: 8 37: action: "NETWORK_ACL_REMOVERULE" options: - - position: 9 + position: 9 38: action: "NETWORK_NIC_DISABLE" options: - - node_id: 6 - - nic_index: 1 + node_id: 6 + nic_id: 1 39: action: "NETWORK_NIC_ENABLE" options: - - node_id: 6 - - nic_index: 1 + node_id: 6 + nic_id: 1 options: nodes: diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 3f674fbb..1e6893ff 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -5,6 +5,8 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from gym import spaces from primaite.simulator.sim_container import Simulation +from primaite import getLogger +_LOGGER = getLogger(__name__) if TYPE_CHECKING: from primaite.game.session import PrimaiteSession @@ -253,7 +255,7 @@ class NodeShutdownAction(NodeAbstractAction): class NodeStartupAction(NodeAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes) - self.verb = "start" + self.verb = "startup" class NodeResetAction(NodeAbstractAction): @@ -274,33 +276,73 @@ class NetworkACLAddRuleAction(AbstractAction): **kwargs, ) -> None: super().__init__(manager=manager) - num_permissions = 2 + num_permissions = 3 self.shape: Dict[str, int] = { "position": max_acl_rules, "permission": num_permissions, - "source_ip_idx": num_ips, - "dest_ip_idx": num_ips, - "source_port_idx": num_ports, - "dest_port_idx": num_ports, - "protocol_idx": num_protocols, + "source_ip_id": num_ips, + "dest_ip_id": num_ips, + "source_port_id": num_ports, + "dest_port_id": num_ports, + "protocol_id": num_protocols, } self.target_router_uuid: str = target_router_uuid def form_request( - self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx + self, position, permission, source_ip_id, dest_ip_id, source_port_id, dest_port_id, protocol_id ) -> List[str]: - protocol = self.manager.get_internet_protocol_by_idx(protocol_idx) - src_ip = self.manager.get_ip_address_by_idx(source_ip_idx) - src_port = self.manager.get_port_by_idx(source_port_idx) - dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx) - dst_port = self.manager.get_port_by_idx(dest_port_idx) + if permission == 0: + permission_str = "UNUSED" + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + elif permission == 1: + permission_str = "ALLOW" + elif permission == 2: + permission_str = "DENY" + else: + _LOGGER.warn(f"{self.__class__} received permission {permission}, expected 0 or 1.") + + if protocol_id == 0: + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + + if protocol_id == 1: + protocol = "ALL" + else: + protocol = self.manager.get_internet_protocol_by_idx(protocol_id-2) + # subtract 2 to account for UNUSED=0 and ALL=1. + + if source_ip_id in [0,1]: + src_ip = "ALL" + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + else: + src_ip = self.manager.get_ip_address_by_idx(source_ip_id-2) + # subtract 2 to account for UNUSED=0, and ALL=1 + + if source_port_id == 1: + src_port = "ALL" + else: + src_port = self.manager.get_port_by_idx(source_port_id-2) + # subtract 2 to account for UNUSED=0, and ALL=1 + + if dest_ip_id in (0,1): + dst_ip = "ALL" + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + else: + dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id) + # subtract 2 to account for UNUSED=0, and ALL=1 + + if dest_port_id == 1: + dst_port = "ALL" + else: + dst_port = self.manager.get_port_by_idx(dest_port_id) + # subtract 2 to account for UNUSED=0, and ALL=1 + return [ "network", "node", self.target_router_uuid, "acl", "add_rule", - permission, + permission_str, protocol, src_ip, src_port, @@ -320,36 +362,52 @@ class NetworkACLRemoveRuleAction(AbstractAction): return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position] -class NetworkNICEnableAction(AbstractAction): +class NetworkNICAbstractAction(AbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} + self.verb: str def form_request(self, node_id: int, nic_id: int) -> List[str]: + node_uuid = self.manager.get_node_uuid_by_idx(node_idx=node_id) + nic_uuid = self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id) + if node_uuid is None or nic_uuid is None: + return ["do_nothing"] return [ "network", "node", - self.manager.get_node_uuid_by_idx(node_idx=node_id), + node_uuid, "nic", - self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), - "enable", + nic_uuid, + self.verb, ] -class NetworkNICDisableAction(AbstractAction): +class NetworkNICEnableAction(NetworkNICAbstractAction): def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} + super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) + self.verb = "enable" - def form_request(self, node_id: int, nic_id: int) -> List[str]: - return [ - "network", - "node", - self.manager.get_node_uuid_by_idx(node_idx=node_id), - "nic", - self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), - "disable", - ] + +class NetworkNICDisableAction(NetworkNICAbstractAction): + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: + super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) + self.verb = "disable" + +# class NetworkNICDisableAction(AbstractAction): +# def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: +# super().__init__(manager=manager) +# self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} + +# def form_request(self, node_id: int, nic_id: int) -> List[str]: +# return [ +# "network", +# "node", +# self.manager.get_node_uuid_by_idx(node_idx=node_id), +# "nic", +# self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), +# "disable", +# ] class ActionManager: diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index c5b931ee..28c87af1 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING +from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING, Sequence, Tuple from gym import spaces from pydantic import BaseModel @@ -15,7 +15,7 @@ the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is """ -def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: +def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: """ Access an item from a deeply dictionary with a list of keys. @@ -29,12 +29,13 @@ def access_from_nested_dict(dictionary: Dict, keys: List[Hashable]) -> Any: :return: The value in the dictionary :rtype: Any """ - if len(keys) == 0: + key_list = [*keys] # copy keys to a new list to prevent editing original list + if len(key_list) == 0: return dictionary - k = keys.pop(0) + k = key_list.pop(0) if k not in dictionary: return NOT_PRESENT_IN_STATE - return access_from_nested_dict(dictionary[k], keys) + return access_from_nested_dict(dictionary[k], key_list) class AbstractObservation(ABC): @@ -66,7 +67,7 @@ class AbstractObservation(ABC): class FileObservation(AbstractObservation): - def __init__(self, where: Optional[List[str]] = None) -> None: + def __init__(self, where: Optional[Tuple[str]] = None) -> None: """ _summary_ @@ -79,7 +80,7 @@ class FileObservation(AbstractObservation): :type where: Optional[List[str]] """ super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where self.default_observation: spaces.Space = {"health_status": 0} "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted." @@ -104,7 +105,7 @@ class ServiceObservation(AbstractObservation): default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} "Default observation is what should be returned when the service doesn't exist." - def __init__(self, where: Optional[List[str]] = None) -> None: + def __init__(self, where: Optional[Tuple[str]] = None) -> None: """ :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with @@ -115,7 +116,7 @@ class ServiceObservation(AbstractObservation): :type where: Optional[List[str]] """ super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: if self.where is None: @@ -124,7 +125,7 @@ class ServiceObservation(AbstractObservation): service_state = access_from_nested_dict(state, self.where) if service_state is NOT_PRESENT_IN_STATE: return self.default_observation - return {"operating_status": service_state["operating_status"], "health_status": service_state["health_status"]} + return {"operating_status": service_state["operating_state"], "health_status": service_state["health_status"]} @property def space(self) -> spaces.Space: @@ -132,7 +133,9 @@ class ServiceObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]=None): - return cls(where=parent_where+["services",session.ref_map_services[config['service_ref']]]) + return cls( + where=parent_where+["services",session.ref_map_services[config['service_ref']].uuid] + ) @@ -140,7 +143,7 @@ class LinkObservation(AbstractObservation): default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} "Default observation is what should be returned when the link doesn't exist." - def __init__(self, where: Optional[List[str]] = None) -> None: + def __init__(self, where: Optional[Tuple[str]] = None) -> None: """ :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with @@ -151,7 +154,7 @@ class LinkObservation(AbstractObservation): :type where: Optional[List[str]] """ super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: if self.where is None: @@ -180,7 +183,7 @@ class LinkObservation(AbstractObservation): class FolderObservation(AbstractObservation): - def __init__(self, where: Optional[List[str]] = None, files: List[FileObservation] = []) -> None: + def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = []) -> None: """Initialise folder Observation, including files inside of the folder. :param where: Where in the simulation state dictionary to find the relevant information for this folder. @@ -199,7 +202,7 @@ class FolderObservation(AbstractObservation): """ super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where self.files: List[FileObservation] = files @@ -246,9 +249,9 @@ class FolderObservation(AbstractObservation): class NicObservation(AbstractObservation): default_observation: spaces.Space = {"nic_status": 0} - def __init__(self, where: Optional[List[str]] = None) -> None: + def __init__(self, where: Optional[Tuple[str]] = None) -> None: super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: if self.where is None: @@ -271,7 +274,7 @@ class NicObservation(AbstractObservation): class NodeObservation(AbstractObservation): def __init__( self, - where: Optional[List[str]] = None, + where: Optional[Tuple[str]] = None, services: List[ServiceObservation] = [], folders: List[FolderObservation] = [], nics: List[NicObservation] = [], @@ -298,7 +301,7 @@ class NodeObservation(AbstractObservation): :type max_nics: int, optional """ super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where self.services: List[ServiceObservation] = services self.folders: List[FolderObservation] = folders @@ -371,10 +374,10 @@ class AclObservation(AbstractObservation): # if a file is created at runtime, we have currently got no way of telling the observation space to track it. # this needs adding, but not for the MVP. def __init__( - self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[List[str]] = None, num_rules: int = 10 + self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[Tuple[str]] = None, num_rules: int = 10 ) -> None: super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where self.num_rules: int = num_rules self.node_to_id: Dict[str, int] = node_ip_to_id "List of node IP addresses, order in this list determines how they are converted to an ID" @@ -403,6 +406,8 @@ class AclObservation(AbstractObservation): if acl_state is NOT_PRESENT_IN_STATE: return self.default_observation + + #TODO: what if the ACL has more rules than num of max rules for obs space obs = {} obs["RULES"] = {} for i, rule_state in acl_state.items(): @@ -466,7 +471,7 @@ class AclObservation(AbstractObservation): node_ip_to_id=node_ip_to_idx, ports=session.options.ports, protocols=session.options.protocols, - where=["network", "nodes", router_uuid]) + where=["network", "nodes", router_uuid, "acl", "acl"]) @@ -498,7 +503,7 @@ class UC2BlueObservation(AbstractObservation): where:Optional[List[str]] = None, ) -> None: super().__init__() - self.where: Optional[List[str]] = where + self.where: Optional[Tuple[str]] = where self.nodes: List[NodeObservation] = nodes self.links: List[LinkObservation] = links @@ -517,11 +522,10 @@ class UC2BlueObservation(AbstractObservation): return self.default_observation obs = {} - obs['NODES'] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} obs['LINKS'] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} - obs['ACL'] = {self.acl.observe(state)} - obs['ICS'] = {self.ics.observe(state)} + obs['ACL'] = self.acl.observe(state) + obs['ICS'] = self.ics.observe(state) return obs @@ -546,7 +550,7 @@ class UC2BlueObservation(AbstractObservation): acl = AclObservation.from_config(config=acl_config, session=session) ics_config = config["ics"] - ics = ICSObservation.from_config(ics_config) + ics = ICSObservation.from_config(config=ics_config, session=session) new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network']) return new diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 2e7681a9..3691c101 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -111,11 +111,11 @@ class AccessControlList(SimComponent): Action( func=lambda request, context: self.add_rule( ACLAction[request[0]], - IPProtocol[request[1]], - IPv4Address[request[2]], - Port[request[3]], - IPv4Address[request[4]], - Port[request[5]], + None if request[1] is "ALL" else IPProtocol[request[1]], + IPv4Address(request[2]), + None if request[3] is "ALL" else Port[request[3]], + IPv4Address(request[4]), + None if request[5] is "ALL" else Port[request[5]], int(request[6]), ) ), diff --git a/src/primaite/simulator/network/hardware/nodes/switch.py b/src/primaite/simulator/network/hardware/nodes/switch.py index ac8dabd1..bb296203 100644 --- a/src/primaite/simulator/network/hardware/nodes/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/switch.py @@ -55,12 +55,11 @@ class Switch(Node): :return: Current state of this object and child objects. """ - return { - "uuid": self.uuid, - "num_ports": self.num_ports, # redundant? - "ports": {port_num: port.describe_state() for port_num, port in self.switch_ports.items()}, - "mac_address_table": {mac: port for mac, port in self.mac_address_table.items()}, - } + state = super().describe_state() + state["ports"] = {port_num: port.describe_state() for port_num, port in self.switch_ports.items()} + state["num_ports"]= self.num_ports # redundant? + state["mac_address_table"]= {mac: port for mac, port in self.mac_address_table.items()} + return state def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort): """ From e85c5977d0de4d644865562e289bafb3a9a1c230 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 9 Oct 2023 18:22:30 +0100 Subject: [PATCH 16/35] remove redundant code from sandbox notebook --- sandbox.ipynb | 792 +------------------------------------------------- 1 file changed, 12 insertions(+), 780 deletions(-) diff --git a/sandbox.ipynb b/sandbox.ipynb index b3b3be0d..a2150921 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -2,568 +2,29 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%load_ext autoreload\n", - "%autoreload 2" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "from primaite.game.session import PrimaiteSession\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ + "%autoreload 2\n", + "from primaite.game.session import PrimaiteSession\n", + "\n", "from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS\n", "import logging\n", "_PRIMAITE_CONFIG['log_level']=logging.DEBUG\n", - "print(PRIMAITE_PATHS.app_log_dir_path)" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "import itertools" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ + "print(PRIMAITE_PATHS.app_log_dir_path)\n", + "import itertools\n", "from primaite.game.session import PrimaiteSession\n", "from primaite.simulator.sim_container import Simulation\n", "from primaite.game.agent.interface import AbstractAgent\n", "from primaite.simulator.network.networks import arcd_uc2_network\n", - "import yaml\n" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [], - "source": [ - "with open('example_config.yaml', 'r') as file:\n", - " cfg = yaml.safe_load(file)" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-08 17:56:35,831: Added node af2f9c15-ecb4-4b65-b48f-63f12acddb88 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,836: Added node 47158854-0917-4037-a6a2-33dde56a120f to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,840: Added node cba8ce63-8064-4f80-bcfe-95ca65221dfa to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,846: Added node e01e7c2b-02ac-4e2d-b7bb-8bc3b6ea6509 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,857: Added node bd5d85ba-5980-45c7-8b28-020a2cfeba0f to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,863: Added node 39e0e37c-4d72-4c76-93cb-4f9c29651ef4 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,878: Added node 7d1063f9-b5e5-4753-966e-1b630325b266 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,881: Added node d85b6abb-0f9e-4853-af26-c9b410e1cb94 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,884: Added node 63b18888-98aa-4182-a014-02999d095bd0 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,888: Added node f514cf8a-a3f1-46d6-be00-994364241ef4 to Network cbd56fbb-104f-4823-9ee6-f4a968343b31\n", - "2023-10-08 17:56:35,890: NIC 27:a9:09:ed:30:5a/192.168.1.1 connected to Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20\n", - "2023-10-08 17:56:35,891: SwitchPort cb:6f:24:8c:7a:20 connected to Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20\n", - "2023-10-08 17:56:35,893: Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20 up\n", - "2023-10-08 17:56:35,896: Link 27:a9:09:ed:30:5a/192.168.1.1<-->cb:6f:24:8c:7a:20 up\n", - "2023-10-08 17:56:35,897: Added link 41d994cb-2976-4aa2-b306-649cef4deb80 to connect 27:a9:09:ed:30:5a/192.168.1.1 and cb:6f:24:8c:7a:20\n", - "2023-10-08 17:56:35,899: NIC 5c:fa:b1:a4:69:ec/192.168.1.1 connected to Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87\n", - "2023-10-08 17:56:35,900: SwitchPort 68:54:d7:42:04:87 connected to Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87\n", - "2023-10-08 17:56:35,901: Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87 up\n", - "2023-10-08 17:56:35,903: Link 5c:fa:b1:a4:69:ec/192.168.1.1<-->68:54:d7:42:04:87 up\n", - "2023-10-08 17:56:35,904: Added link d582a248-e968-40eb-9d1b-67143d729e0c to connect 5c:fa:b1:a4:69:ec/192.168.1.1 and 68:54:d7:42:04:87\n", - "2023-10-08 17:56:35,905: SwitchPort c6:bd:77:78:4b:5d connected to Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10\n", - "2023-10-08 17:56:35,908: Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10 up\n", - "2023-10-08 17:56:35,909: NIC 1c:d9:92:e8:d6:3b/192.168.1.10 connected to Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10\n", - "2023-10-08 17:56:35,911: Link c6:bd:77:78:4b:5d<-->1c:d9:92:e8:d6:3b/192.168.1.10 up\n", - "2023-10-08 17:56:35,912: Added link 13315780-1fcc-4c85-b94b-ef8f14c88a8a to connect c6:bd:77:78:4b:5d and 1c:d9:92:e8:d6:3b/192.168.1.10\n", - "2023-10-08 17:56:35,913: SwitchPort cd:46:af:c4:33:65 connected to Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12\n", - "2023-10-08 17:56:35,916: Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12 up\n", - "2023-10-08 17:56:35,917: NIC aa:cf:2f:71:13:5b/192.168.1.12 connected to Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12\n", - "2023-10-08 17:56:35,918: Link cd:46:af:c4:33:65<-->aa:cf:2f:71:13:5b/192.168.1.12 up\n", - "2023-10-08 17:56:35,919: Added link 6c2a80f7-f36d-4df6-ac84-d354e5d517dd to connect cd:46:af:c4:33:65 and aa:cf:2f:71:13:5b/192.168.1.12\n", - "2023-10-08 17:56:35,920: SwitchPort 2c:d2:67:ef:68:a8 connected to Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14\n", - "2023-10-08 17:56:35,923: Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14 up\n", - "2023-10-08 17:56:35,924: NIC e1:09:5e:98:ee:a2/192.168.1.14 connected to Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14\n", - "2023-10-08 17:56:35,925: Link 2c:d2:67:ef:68:a8<-->e1:09:5e:98:ee:a2/192.168.1.14 up\n", - "2023-10-08 17:56:35,926: Added link 1cfdd4f2-22be-4e69-8f1f-daef8e18f543 to connect 2c:d2:67:ef:68:a8 and e1:09:5e:98:ee:a2/192.168.1.14\n", - "2023-10-08 17:56:35,927: SwitchPort 9b:13:8c:a0:8c:82 connected to Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16\n", - "2023-10-08 17:56:35,929: Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16 up\n", - "2023-10-08 17:56:35,930: NIC cc:c2:84:03:1c:42/192.168.1.16 connected to Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16\n", - "2023-10-08 17:56:35,932: Link 9b:13:8c:a0:8c:82<-->cc:c2:84:03:1c:42/192.168.1.16 up\n", - "2023-10-08 17:56:35,933: Added link 031111e1-3b05-49ce-bd1f-2cdf77b210f4 to connect 9b:13:8c:a0:8c:82 and cc:c2:84:03:1c:42/192.168.1.16\n", - "2023-10-08 17:56:35,934: SwitchPort a1:70:9e:43:1c:07 connected to Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110\n", - "2023-10-08 17:56:35,937: Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110 up\n", - "2023-10-08 17:56:35,938: NIC e7:58:3c:ed:f7:37/192.168.1.110 connected to Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110\n", - "2023-10-08 17:56:35,939: Link a1:70:9e:43:1c:07<-->e7:58:3c:ed:f7:37/192.168.1.110 up\n", - "2023-10-08 17:56:35,941: Added link f15884e7-0df6-4fa5-bb72-406cb2bdff45 to connect a1:70:9e:43:1c:07 and e7:58:3c:ed:f7:37/192.168.1.110\n", - "2023-10-08 17:56:35,943: SwitchPort a5:da:f2:03:21:e3 connected to Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21\n", - "2023-10-08 17:56:35,946: Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21 up\n", - "2023-10-08 17:56:35,947: NIC cf:63:9f:62:fe:df/192.168.10.21 connected to Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21\n", - "2023-10-08 17:56:35,948: Link a5:da:f2:03:21:e3<-->cf:63:9f:62:fe:df/192.168.10.21 up\n", - "2023-10-08 17:56:35,950: Added link cc6767fa-25de-4daa-bd47-37b49b15a881 to connect a5:da:f2:03:21:e3 and cf:63:9f:62:fe:df/192.168.10.21\n", - "2023-10-08 17:56:35,951: SwitchPort eb:5b:86:14:bd:d1 connected to Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22\n", - "2023-10-08 17:56:35,953: Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22 up\n", - "2023-10-08 17:56:35,954: NIC 3d:73:a6:62:97:3a/192.168.10.22 connected to Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22\n", - "2023-10-08 17:56:35,955: Link eb:5b:86:14:bd:d1<-->3d:73:a6:62:97:3a/192.168.10.22 up\n", - "2023-10-08 17:56:35,958: Added link b29563ed-0636-4188-ba28-52b74a04da27 to connect eb:5b:86:14:bd:d1 and 3d:73:a6:62:97:3a/192.168.10.22\n", - "2023-10-08 17:56:35,959: SwitchPort e3:3a:0b:03:0b:8c connected to Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110\n", - "2023-10-08 17:56:35,961: Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110 up\n", - "2023-10-08 17:56:35,963: NIC d4:ff:37:8d:e4:3d/192.168.10.110 connected to Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110\n", - "2023-10-08 17:56:35,963: Link e3:3a:0b:03:0b:8c<-->d4:ff:37:8d:e4:3d/192.168.10.110 up\n", - "2023-10-08 17:56:35,964: Added link f6fe7757-a8c1-4cdc-a2b2-49d247117903 to connect e3:3a:0b:03:0b:8c and d4:ff:37:8d:e4:3d/192.168.10.110\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "service type not found DatabaseBackup\n", - "service type not found WebBrowser\n" - ] - } - ], - "source": [ - "sess = PrimaiteSession.from_config(cfg)" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "[,\n", - " ,\n", - " ]" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sess.agents" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [], - "source": [ - "network = sess.simulation.network" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-08 17:56:36,041: Stepping primaite session. Step counter: 0\n", - "2023-10-08 17:56:36,043: Sending simulation state to agent client_1_green_user\n", - "2023-10-08 17:56:36,045: Getting agent action\n", - "2023-10-08 17:56:36,047: Formatting agent action DONOTHING\n", - "2023-10-08 17:56:36,048: Sending request to simulation: ['do_nothing']\n", - "2023-10-08 17:56:36,050: Sending simulation state to agent client_1_data_manipulation_red_bot\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[]\n", - "[]\n" - ] - }, - { - "ename": "TypeError", - "evalue": "unhashable type: 'DataManipulationBot'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", - "\u001b[1;32m/home/cade/repos/PrimAITE/sandbox.ipynb Cell 10\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> 1\u001b[0m sess\u001b[39m.\u001b[39;49mstep()\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/session.py:80\u001b[0m, in \u001b[0;36mPrimaiteSession.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 77\u001b[0m sim_state \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msimulation\u001b[39m.\u001b[39mdescribe_state()\n\u001b[1;32m 79\u001b[0m \u001b[39m# 6. each agent takes most recent state and converts it to CAOS observation\u001b[39;00m\n\u001b[0;32m---> 80\u001b[0m agent_obs \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mconvert_state_to_obs(sim_state)\n\u001b[1;32m 82\u001b[0m \u001b[39m# 7. meanwhile each agent also takes state and calculates reward\u001b[39;00m\n\u001b[1;32m 83\u001b[0m agent_reward \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39mcalculate_reward_from_state(sim_state)\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/interface.py:40\u001b[0m, in \u001b[0;36mAbstractAgent.convert_state_to_obs\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mconvert_state_to_obs\u001b[39m(\u001b[39mself\u001b[39m, state: Dict) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m ObsType:\n\u001b[1;32m 36\u001b[0m \u001b[39m \u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 37\u001b[0m \u001b[39m state : dict state directly from simulation.describe_state\u001b[39;00m\n\u001b[1;32m 38\u001b[0m \u001b[39m output : dict state according to CAOS.\u001b[39;00m\n\u001b[1;32m 39\u001b[0m \u001b[39m \"\"\"\u001b[39;00m\n\u001b[0;32m---> 40\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mobservation_space\u001b[39m.\u001b[39;49mobserve(state)\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:608\u001b[0m, in \u001b[0;36mObservationSpace.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 607\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mobserve\u001b[39m(\u001b[39mself\u001b[39m, state) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Dict:\n\u001b[0;32m--> 608\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mobs\u001b[39m.\u001b[39;49mobserve(state)\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:571\u001b[0m, in \u001b[0;36mUC2RedObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 568\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[1;32m 570\u001b[0m obs \u001b[39m=\u001b[39m {}\n\u001b[0;32m--> 571\u001b[0m obs[\u001b[39m'\u001b[39m\u001b[39mNODES\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m {i\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m: node\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, node \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes)}\n\u001b[1;32m 572\u001b[0m \u001b[39mreturn\u001b[39;00m obs\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:571\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 568\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[1;32m 570\u001b[0m obs \u001b[39m=\u001b[39m {}\n\u001b[0;32m--> 571\u001b[0m obs[\u001b[39m'\u001b[39m\u001b[39mNODES\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m {i\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m: node\u001b[39m.\u001b[39;49mobserve(state) \u001b[39mfor\u001b[39;00m i, node \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnodes)}\n\u001b[1;32m 572\u001b[0m \u001b[39mreturn\u001b[39;00m obs\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:328\u001b[0m, in \u001b[0;36mNodeObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 326\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)\n\u001b[1;32m 327\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)\n\u001b[0;32m--> 328\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mSERVICES\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: service\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, service \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)}\n\u001b[1;32m 329\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mFOLDERS\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: folder\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, folder \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)}\n\u001b[1;32m 330\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39moperating_status\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m node_state[\u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m]\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:328\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 326\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)\n\u001b[1;32m 327\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)\n\u001b[0;32m--> 328\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mSERVICES\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: service\u001b[39m.\u001b[39;49mobserve(state) \u001b[39mfor\u001b[39;00m i, service \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mservices)}\n\u001b[1;32m 329\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39mFOLDERS\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m {i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m: folder\u001b[39m.\u001b[39mobserve(state) \u001b[39mfor\u001b[39;00m i, folder \u001b[39min\u001b[39;00m \u001b[39menumerate\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfolders)}\n\u001b[1;32m 330\u001b[0m obs[\u001b[39m\"\u001b[39m\u001b[39moperating_status\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m node_state[\u001b[39m\"\u001b[39m\u001b[39moperating_state\u001b[39m\u001b[39m\"\u001b[39m]\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:124\u001b[0m, in \u001b[0;36mServiceObservation.observe\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m 121\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mwhere \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 122\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n\u001b[0;32m--> 124\u001b[0m service_state \u001b[39m=\u001b[39m access_from_nested_dict(state, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mwhere)\n\u001b[1;32m 125\u001b[0m \u001b[39mif\u001b[39;00m service_state \u001b[39mis\u001b[39;00m NOT_PRESENT_IN_STATE:\n\u001b[1;32m 126\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdefault_observation\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", - " \u001b[0;31m[... skipping similar frames: access_from_nested_dict at line 37 (1 times)]\u001b[0m\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:37\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[0;32m---> 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", - "File \u001b[0;32m~/repos/PrimAITE/src/primaite/game/agent/observations.py:35\u001b[0m, in \u001b[0;36maccess_from_nested_dict\u001b[0;34m(dictionary, keys)\u001b[0m\n\u001b[1;32m 33\u001b[0m \u001b[39mreturn\u001b[39;00m dictionary\n\u001b[1;32m 34\u001b[0m k \u001b[39m=\u001b[39m keys\u001b[39m.\u001b[39mpop(\u001b[39m0\u001b[39m)\n\u001b[0;32m---> 35\u001b[0m \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;49;00m \u001b[39min\u001b[39;49;00m dictionary:\n\u001b[1;32m 36\u001b[0m \u001b[39mreturn\u001b[39;00m NOT_PRESENT_IN_STATE\n\u001b[1;32m 37\u001b[0m \u001b[39mreturn\u001b[39;00m access_from_nested_dict(dictionary[k], keys)\n", - "\u001b[0;31mTypeError\u001b[0m: unhashable type: 'DataManipulationBot'" - ] - } - ], - "source": [ - "sess.step()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from ipaddress import IPv4Address\n", - "\n", - "from primaite.simulator.network.container import Network\n", - "from primaite.simulator.network.hardware.base import NIC\n", - "from primaite.simulator.network.hardware.nodes.computer import Computer\n", - "from primaite.simulator.network.hardware.nodes.router import ACLAction, Router\n", - "from primaite.simulator.network.hardware.nodes.server import Server\n", - "from primaite.simulator.network.hardware.nodes.switch import Switch\n", - "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", - "from primaite.simulator.network.transmission.transport_layer import Port\n", - "from primaite.simulator.system.applications.database_client import DatabaseClient\n", - "from primaite.simulator.system.services.database_service import DatabaseService\n", - "from primaite.simulator.system.services.dns_client import DNSClient\n", - "from primaite.simulator.system.services.dns_server import DNSServer\n", - "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "router_1 = Router(hostname=\"router_1\", num_ports=5)\n", - "router_1.power_on()\n", - "router_1.configure_port(port=1, ip_address=\"192.168.1.1\", subnet_mask=\"255.255.255.0\")\n", - "router_1.configure_port(port=2, ip_address=\"192.168.10.1\", subnet_mask=\"255.255.255.0\")\n", - "\n", - "# Switch 1\n", - "switch_1 = Switch(hostname=\"switch_1\", num_ports=8)\n", - "switch_1.power_on()\n", - "network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[8])\n", - "router_1.enable_port(1)\n", - "\n", - "# Switch 2\n", - "switch_2 = Switch(hostname=\"switch_2\", num_ports=8)\n", - "switch_2.power_on()\n", - "network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[8])\n", - "router_1.enable_port(2)\n", - "\n", - "# Client 1\n", - "client_1 = Computer(\n", - " hostname=\"client_1\",\n", - " ip_address=\"192.168.10.21\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.10.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "client_1.power_on()\n", - "client_1.software_manager.install(DNSClient)\n", - "client_1_dns_client_service: DNSServer = client_1.software_manager.software[\"DNSClient\"] # noqa\n", - "client_1_dns_client_service.start()\n", - "network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])\n", - "client_1.software_manager.install(DataManipulationBot)\n", - "db_manipulation_bot: DataManipulationBot = client_1.software_manager.software[\"DataManipulationBot\"]\n", - "db_manipulation_bot.configure(server_ip_address=IPv4Address(\"192.168.1.14\"), payload=\"DROP TABLE IF EXISTS user;\")\n", - "\n", - "# Client 2\n", - "client_2 = Computer(\n", - " hostname=\"client_2\",\n", - " ip_address=\"192.168.10.22\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.10.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "client_2.power_on()\n", - "client_2.software_manager.install(DNSClient)\n", - "client_2_dns_client_service: DNSServer = client_2.software_manager.software[\"DNSClient\"] # noqa\n", - "client_2_dns_client_service.start()\n", - "network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2])\n", - "\n", - "# Domain Controller\n", - "domain_controller = Server(\n", - " hostname=\"domain_controller\",\n", - " ip_address=\"192.168.1.10\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.1.1\",\n", - ")\n", - "domain_controller.power_on()\n", - "domain_controller.software_manager.install(DNSServer)\n", - "\n", - "network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1])\n", - "\n", - "# Database Server\n", - "database_server = Server(\n", - " hostname=\"database_server\",\n", - " ip_address=\"192.168.1.14\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.1.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "database_server.power_on()\n", - "network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3])\n", - "\n", - "ddl = \"\"\"\n", - "CREATE TABLE IF NOT EXISTS user (\n", - "id INTEGER PRIMARY KEY AUTOINCREMENT,\n", - "name VARCHAR(50) NOT NULL,\n", - "email VARCHAR(50) NOT NULL,\n", - "age INT,\n", - "city VARCHAR(50),\n", - "occupation VARCHAR(50)\n", - ");\"\"\"\n", - "\n", - "user_insert_statements = [\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('John Doe', 'johndoe@example.com', 32, 'New York', 'Engineer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Jane Smith', 'janesmith@example.com', 27, 'Los Angeles', 'Designer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Bob Johnson', 'bobjohnson@example.com', 45, 'Chicago', 'Manager');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Alice Lee', 'alicelee@example.com', 22, 'San Francisco', 'Student');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('David Kim', 'davidkim@example.com', 38, 'Houston', 'Consultant');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Emily Chen', 'emilychen@example.com', 29, 'Seattle', 'Software Developer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Frank Wang', 'frankwang@example.com', 55, 'New York', 'Entrepreneur');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Grace Park', 'gracepark@example.com', 31, 'Los Angeles', 'Marketing Specialist');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Henry Wu', 'henrywu@example.com', 40, 'Chicago', 'Accountant');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Isabella Kim', 'isabellakim@example.com', 26, 'San Francisco', 'Graphic Designer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Jake Lee', 'jakelee@example.com', 33, 'Houston', 'Sales Manager');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Kelly Chen', 'kellychen@example.com', 28, 'Seattle', 'Web Developer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Lucas Liu', 'lucasliu@example.com', 42, 'New York', 'Lawyer');\", # noqa\n", - " \"INSERT INTO user (name, email, age, city, occupation) VALUES ('Maggie Wang', 'maggiewang@example.com', 30, 'Los Angeles', 'Data Analyst');\", # noqa\n", - "]\n", - "database_server.software_manager.install(DatabaseService)\n", - "database_service: DatabaseService = database_server.software_manager.software[\"DatabaseService\"] # noqa\n", - "database_service.start()\n", - "database_service._process_sql(ddl, None) # noqa\n", - "for insert_statement in user_insert_statements:\n", - " database_service._process_sql(insert_statement, None) # noqa\n", - "\n", - "# Web Server\n", - "web_server = Server(\n", - " hostname=\"web_server\",\n", - " ip_address=\"192.168.1.12\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.1.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "web_server.power_on()\n", - "web_server.software_manager.install(DatabaseClient)\n", - "\n", - "database_client: DatabaseClient = web_server.software_manager.software[\"DatabaseClient\"]\n", - "database_client.configure(server_ip_address=IPv4Address(\"192.168.1.14\"))\n", - "network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])\n", - "database_client.run()\n", - "database_client.connect()\n", - "\n", - "# register the web_server to a domain\n", - "dns_server_service: DNSServer = domain_controller.software_manager.software[\"DNSServer\"] # noqa\n", - "dns_server_service.start()\n", - "dns_server_service.dns_register(\"arcd.com\", web_server.ip_address)\n", - "\n", - "# Backup Server\n", - "backup_server = Server(\n", - " hostname=\"backup_server\",\n", - " ip_address=\"192.168.1.16\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.1.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "backup_server.power_on()\n", - "network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4])\n", - "\n", - "# Security Suite\n", - "security_suite = Server(\n", - " hostname=\"security_suite\",\n", - " ip_address=\"192.168.1.110\",\n", - " subnet_mask=\"255.255.255.0\",\n", - " default_gateway=\"192.168.1.1\",\n", - " dns_server=IPv4Address(\"192.168.1.10\"),\n", - ")\n", - "security_suite.power_on()\n", - "network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7])\n", - "security_suite.connect_nic(NIC(ip_address=\"192.168.10.110\", subnet_mask=\"255.255.255.0\"))\n", - "network.connect(endpoint_b=security_suite.ethernet_port[2], endpoint_a=switch_2.switch_ports[7])\n", - "\n", - "router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)\n", - "\n", - "router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)\n", - "\n", - "# Allow PostgreSQL requests\n", - "router_1.acl.add_rule(\n", - " action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0\n", - ")\n", - "\n", - "# Allow DNS requests\n", - "router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "node_uuid_list = list(sess.simulation.network.nodes.keys())" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from primaite.game.agent.actions import ActionManager" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "actman = ActionManager(sess.simulation, [\"DONOTHING\", \"NODE_SERVICE_SCAN\", \"NODE_SERVICE_STOP\", \"NODE_FOLDER_SCAN\"],node_uuid_list,act_map={\n", - " 0:{\n", - " \"action\": \"DONOTHING\",\n", - " \"options\": {}\n", - " },\n", - " 1:{\n", - " \"action\": \"NODE_SERVICE_SCAN\",\n", - " \"options\": {\"node_id\":0, \"service_id\":0},\n", - " },\n", - " 2:{\n", - " \"action\": \"NODE_SERVICE_SCAN\",\n", - " \"options\": {\"node_id\":1, \"service_id\":0},\n", - " },\n", - " 3:{\n", - " \"action\": \"NODE_FOLDER_SCAN\",\n", - " \"options\": {\"node_id\":4, \"folder_id\":0},\n", - " }\n", - "})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "act_id, act_options = actman.get_action(3)\n", - "my_trial_act = actman.form_request(action_identifier=act_id, action_options=act_options)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sess.simulation.apply_action(my_trial_act)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "my_trial_act" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sess.step()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sess.step_counter" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from gym import spaces" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sp = spaces.Tuple( (spaces.MultiDiscrete([3, 2]), spaces.MultiDiscrete([3, 2]), spaces.MultiDiscrete([3, 2]),))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sp.sample()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ "import yaml\n", - "from primaite.simulator.sim_container import Simulation\n", - "from primaite.simulator.network.hardware.nodes.computer import Computer\n", - "from primaite.simulator.network.hardware.nodes.server import Server\n", - "from primaite.simulator.network.hardware.nodes.switch import Switch\n", - "from primaite.simulator.network.hardware.nodes.router import Router\n", "\n", - "from primaite.simulator.system.applications.database_client import DatabaseClient\n", - "from primaite.simulator.system.services.database_service import DatabaseService\n", - "from primaite.simulator.system.services.dns_client import DNSClient\n", - "from primaite.simulator.system.services.dns_server import DNSServer\n", - "from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot\n", - "\n", - "\n", - "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", - "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", - "from primaite.simulator.network.transmission.transport_layer import Port\n", - "\n", - "from ipaddress import IPv4Address\n" + "with open('example_config.yaml', 'r') as file:\n", + " cfg = yaml.safe_load(file)\n", + "sess = PrimaiteSession.from_config(cfg)\n", + "sess.step()" ] }, { @@ -572,238 +33,9 @@ "metadata": {}, "outputs": [], "source": [ - "# import yaml\n", - "\n", - "\n", - "from typing import Dict\n", - "from primaite.game.agent.interface import AbstractAgent\n", - "from primaite.game.agent.observations import AclObservation, FileObservation, FolderObservation, ICSObservation, LinkObservation, NicObservation, NodeObservation, NullObservation, ServiceObservation, UC2BlueObservation, UC2RedObservation\n", - "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", - "from primaite.simulator.system.services.service import Service\n", - "\n", - "from primaite.game.agent.scripted_agents import GreenWebBrowsingAgent, RedDatabaseCorruptingAgent\n", - "from primaite.game.agent.GATE_agents import GATERLAgent\n", - "\n", - "class PrimaiteSession:\n", - "\n", - " def __init__(self):\n", - " self.simulation: Simulation\n", - " self.agents = []\n", - "\n", - " @classmethod\n", - " def from_config(cls, cfg_path):\n", - " ref_map_nodes: Dict[str,Node] = {}\n", - " ref_map_services: Dict[str, Service] = {}\n", - " ref_map_links: Dict[str, Link] = {}\n", - " # ref_map_agents: Dict[str, AgentInterface] = {}\n", - "\n", - "\n", - " session = cls()\n", - " with open(cfg_path, 'r') as file:\n", - " conf = yaml.safe_load(file)\n", - "\n", - " #1. create nodes\n", - " sim = Simulation()\n", - " net = sim.network\n", - " nodes_cfg = conf['simulation']['network']['nodes']\n", - " links_cfg = conf['simulation']['network']['links']\n", - " for node_cfg in nodes_cfg:\n", - " node_ref = node_cfg['ref']\n", - " n_type = node_cfg['type']\n", - " if n_type == 'computer':\n", - " new_node = Computer(hostname = node_cfg['hostname'],\n", - " ip_address = node_cfg['ip_address'],\n", - " subnet_mask = node_cfg['subnet_mask'],\n", - " default_gateway = node_cfg['default_gateway'],\n", - " dns_server = node_cfg['dns_server'])\n", - " elif n_type == 'server':\n", - " new_node = Server(hostname = node_cfg['hostname'],\n", - " ip_address = node_cfg['ip_address'],\n", - " subnet_mask = node_cfg['subnet_mask'],\n", - " default_gateway = node_cfg['default_gateway'],\n", - " dns_server = node_cfg.get('dns_server'))\n", - " elif n_type == 'switch':\n", - " new_node = Switch(hostname = node_cfg['hostname'],\n", - " num_ports = node_cfg.get('num_ports'))\n", - " elif n_type == 'router':\n", - " new_node = Router(hostname=node_cfg['hostname'],\n", - " num_ports = node_cfg.get('num_ports'))\n", - " if 'ports' in node_cfg:\n", - " for port_num, port_cfg in node_cfg['ports'].items():\n", - " new_node.configure_port(port=port_num,\n", - " ip_address=port_cfg['ip_address'],\n", - " subnet_mask=port_cfg['subnet_mask'])\n", - " if 'acl' in node_cfg:\n", - " for r_num, r_cfg in node_cfg['acl'].items():\n", - " # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating\n", - " # this: 'r_cfg.get('src_port')'\n", - " # Port/IPProtocol. TODO Refactor\n", - " new_node.acl.add_rule(\n", - " action = ACLAction[r_cfg['action']],\n", - " src_port = None if not (p:=r_cfg.get('src_port')) else Port[p],\n", - " dst_port = None if not (p:=r_cfg.get('dst_port')) else Port[p],\n", - " protocol = None if not (p:=r_cfg.get('protocol')) else IPProtocol[p],\n", - " src_ip_address = r_cfg.get('ip_address'),\n", - " dst_ip_address = r_cfg.get('ip_address'),\n", - " position = r_num\n", - " )\n", - " else:\n", - " print('invalid node type')\n", - " if 'services' in node_cfg:\n", - " for service_cfg in node_cfg['services']:\n", - " service_ref = service_cfg['ref']\n", - " service_type = service_cfg['type']\n", - " service_types_mapping = {\n", - " 'DNSClient': DNSClient, # key is equal to the 'name' attr of the service class itself.\n", - " 'DNSServer' : DNSServer,\n", - " 'DatabaseClient': DatabaseClient,\n", - " 'DatabaseService': DatabaseService,\n", - " # 'database_backup': ,\n", - " 'DataManipulationBot': DataManipulationBot,\n", - " # 'web_browser'\n", - " }\n", - " if service_type in service_types_mapping:\n", - " new_node.software_manager.install(service_types_mapping[service_type])\n", - " new_service = new_node.software_manager.software[service_type]\n", - " ref_map_services[service_ref] = new_service\n", - " else:\n", - " print(f\"service type not found {service_type}\")\n", - " # service-dependent options\n", - " if service_type == 'DatabaseClient':\n", - " if 'options' in service_cfg:\n", - " opt = service_cfg['options']\n", - " if 'db_server_ip' in opt:\n", - " new_service.configure(server_ip_address=IPv4Address(opt['db_server_ip']))\n", - " if service_type == 'DNSServer':\n", - " if 'options' in service_cfg:\n", - " opt = service_cfg['options']\n", - " if 'domain_mapping' in opt:\n", - " for domain, ip in opt['domain_mapping'].items():\n", - " new_service.dns_register(domain, ip)\n", - " if 'nics' in node_cfg:\n", - " for nic_num, nic_cfg in node_cfg['nics'].items():\n", - " new_node.connect_nic(NIC(ip_address=nic_cfg['ip_address'], subnet_mask=nic_cfg['subnet_mask']))\n", - "\n", - " net.add_node(new_node)\n", - " new_node.power_on()\n", - " ref_map_nodes[node_ref] = new_node.uuid\n", - "\n", - " #2. create links between nodes\n", - " for link_cfg in links_cfg:\n", - " node_a = net.nodes[ref_map_nodes[link_cfg['endpoint_a_ref']]]\n", - " node_b = net.nodes[ref_map_nodes[link_cfg['endpoint_b_ref']]]\n", - " if isinstance(node_a, Switch):\n", - " endpoint_a = node_a.switch_ports[link_cfg['endpoint_a_port']]\n", - " else:\n", - " endpoint_a = node_a.ethernet_port[link_cfg['endpoint_a_port']]\n", - " if isinstance(node_b, Switch):\n", - " endpoint_b = node_b.switch_ports[link_cfg['endpoint_b_port']]\n", - " else:\n", - " endpoint_b = node_b.ethernet_port[link_cfg['endpoint_b_port']]\n", - " new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)\n", - " ref_map_links[link_cfg['ref']] = new_link.uuid\n", - "\n", - " session.simulation = sim\n", - " #3. create agents\n", - " game_cfg = conf['game_config']\n", - " ports_cfg = game_cfg['ports']\n", - " protocols_cfg = game_cfg['protocols']\n", - " agents_cfg = game_cfg['agents']\n", - "\n", - " for agent_cfg in agents_cfg:\n", - " agent_ref = agent_cfg['ref']\n", - " agent_type = agent_cfg['type']\n", - " action_space_cfg = agent_cfg['action_space']\n", - " observation_space_cfg = agent_cfg['observation_space']\n", - " reward_function_cfg = agent_cfg['reward_function']\n", - "\n", - " # CREATE OBSERVATION SPACE\n", - " if observation_space_cfg is None:\n", - " obs_space = NullObservation()\n", - " elif observation_space_cfg['type'] == 'UC2BlueObservation':\n", - " node_obs_list = []\n", - " link_obs_list = []\n", - "\n", - "\n", - " #node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses\n", - " node_ip_to_index = {}\n", - " for node_idx, node_cfg in enumerate(nodes_cfg):\n", - " n_ref = node_cfg['ref']\n", - " n_obj = net.nodes[ref_map_nodes[n_ref]]\n", - " for nic_uuid, nic_obj in n_obj.nics.items():\n", - " node_ip_to_index[nic_obj.ip_address] = node_idx + 2\n", - "\n", - "\n", - "\n", - " for node_obs_cfg in observation_space_cfg['options']['nodes']:\n", - " node_ref = node_obs_cfg['node_ref']\n", - " folder_obs_list = []\n", - " service_obs_list = []\n", - " if 'services' in node_obs_cfg:\n", - " for service_obs_cfg in node_obs_cfg['services']:\n", - " service_obs_list.append(ServiceObservation(where=['network','nodes',ref_map_nodes[node_ref],'services',ref_map_services[service_obs_cfg['service_ref']]]))\n", - " if 'folders' in node_obs_cfg:\n", - " for folder_obs_cfg in node_obs_cfg['folders']:\n", - " file_obs_list = []\n", - " if 'files' in folder_obs_cfg:\n", - " for file_obs_cfg in folder_obs_cfg['files']:\n", - " file_obs_list.append(FileObservation(where=['network','nodes',ref_map_nodes[node_ref], 'folders',folder_obs_cfg['folder_name'], 'files', file_obs_cfg['file_name']]))\n", - " folder_obs_list.append(FolderObservation(where=['network','nodes',ref_map_nodes[node_ref], 'folders',folder_obs_cfg['folder_name']], files=file_obs_list))\n", - " nic_obs_list = []\n", - " for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg['node_ref']]].nics.keys():\n", - " nic_obs_list.append(NicObservation(where=['network','nodes',ref_map_nodes[node_ref],'NICs',nic_uuid]))\n", - " node_obs_list.append(NodeObservation(where=['network','nodes',ref_map_nodes[node_ref]], services=service_obs_list, folders=folder_obs_list,nics=nic_obs_list, logon_status=False))\n", - " for link_obs_cfg in observation_space_cfg['options']['links']:\n", - " link_ref = link_obs_cfg['link_ref']\n", - " link_obs_list.append(LinkObservation(where=['network' ,'links', ref_map_links[link_ref]]))\n", - "\n", - " acl_obs = AclObservation(node_ip_to_id=node_ip_to_index, ports=game_cfg['ports'], protocols=game_cfg['ports'], where=['network','nodes',observation_space_cfg['options']['acl']['router_node_ref']])\n", - " obs_space = UC2BlueObservation(nodes=node_obs_list,links=link_obs_list,acl=acl_obs, ics=ICSObservation())\n", - " elif observation_space_cfg['type'] == 'UC2RedObservation':\n", - " obs_space = UC2RedObservation.from_config(observation_space_cfg['options'], sim=sim)\n", - " else:\n", - " print(\"observation space config not specified correctly.\")\n", - " obs_space = NullObservation()\n", - "\n", - " # CREATE ACTION SPACE\n", - "\n", - "\n", - "\n", - " # CREATE REWARD FUNCTION\n", - "\n", - " # CREATE AGENT\n", - " if agent_type == 'GreenWebBrowsingAgent':\n", - " ...\n", - " elif agent_type == 'GATERLAgent':\n", - " ...\n", - " elif agent_type == 'RedDatabaseCorruptingAgent':\n", - " ...\n", - " else:\n", - " print(\"agent type not found\")\n", - "\n", - "\n", - " #4. set up agents' actions and observation spaces.\n", - " return session\n", - "\n", - "s = PrimaiteSession.from_config('example_config.yaml')\n", - "# print(s.simulation.describe_state())" + "for i in range(50):\n", + " sess.step()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s.agents" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { From c9bc8fbf3d512f2ba3fecec1ff008ad2e8a57bb3 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 9 Oct 2023 18:33:30 +0100 Subject: [PATCH 17/35] Fix file observation test --- tests/integration_tests/game_layer/test_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index 7f20a938..c1f20d78 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -16,5 +16,5 @@ def test_file_observation(): dog_file_obs = FileObservation( where=["network", "nodes", pc.uuid, "file_system", "folders", "root", "files", "dog.png"] ) - assert dog_file_obs(state) == {"health_status": 1} + assert dog_file_obs.observe(state) == {"health_status": 1} assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)}) From 91f06c15f6d52ef821dde93f5ad7ec7d23f95072 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 9 Oct 2023 18:35:30 +0100 Subject: [PATCH 18/35] Fix formatting with precommit --- src/primaite/game/agent/actions.py | 67 +++---- src/primaite/game/agent/interface.py | 6 +- src/primaite/game/agent/observations.py | 185 ++++++++++-------- src/primaite/game/agent/rewards.py | 26 +-- src/primaite/game/session.py | 61 +++--- .../network/hardware/nodes/switch.py | 4 +- 6 files changed, 189 insertions(+), 160 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 1e6893ff..cba90305 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -4,8 +4,9 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from gym import spaces -from primaite.simulator.sim_container import Simulation from primaite import getLogger +from primaite.simulator.sim_container import Simulation + _LOGGER = getLogger(__name__) if TYPE_CHECKING: @@ -90,56 +91,56 @@ class NodeServiceAbstractAction(AbstractAction): class NodeServiceScanAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "scan" class NodeServiceStopAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "stop" class NodeServiceStartAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "start" class NodeServicePauseAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "pause" class NodeServiceResumeAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "resume" class NodeServiceRestartAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "restart" class NodeServiceDisableAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "disable" class NodeServiceEnableAction(NodeServiceAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_services:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "enable" class NodeFolderAbstractAction(AbstractAction): @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders} self.verb: str @@ -153,25 +154,25 @@ class NodeFolderAbstractAction(AbstractAction): class NodeFolderScanAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "scan" class NodeFolderCheckhashAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "checkhash" class NodeFolderRepairAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "repair" class NodeFolderRestoreAction(NodeFolderAbstractAction): - def __init__(self, manager: "ActionManager", num_nodes:int, num_folders:int, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "restore" @@ -293,7 +294,7 @@ class NetworkACLAddRuleAction(AbstractAction): ) -> List[str]: if permission == 0: permission_str = "UNUSED" - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS elif permission == 1: permission_str = "ALLOW" elif permission == 2: @@ -302,30 +303,30 @@ class NetworkACLAddRuleAction(AbstractAction): _LOGGER.warn(f"{self.__class__} received permission {permission}, expected 0 or 1.") if protocol_id == 0: - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS if protocol_id == 1: protocol = "ALL" else: - protocol = self.manager.get_internet_protocol_by_idx(protocol_id-2) + protocol = self.manager.get_internet_protocol_by_idx(protocol_id - 2) # subtract 2 to account for UNUSED=0 and ALL=1. - if source_ip_id in [0,1]: + if source_ip_id in [0, 1]: src_ip = "ALL" - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS else: - src_ip = self.manager.get_ip_address_by_idx(source_ip_id-2) + src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 if source_port_id == 1: src_port = "ALL" else: - src_port = self.manager.get_port_by_idx(source_port_id-2) + src_port = self.manager.get_port_by_idx(source_port_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 - if dest_ip_id in (0,1): + if dest_ip_id in (0, 1): dst_ip = "ALL" - return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS + return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS else: dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id) # subtract 2 to account for UNUSED=0, and ALL=1 @@ -394,6 +395,7 @@ class NetworkNICDisableAction(NetworkNICAbstractAction): super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) self.verb = "disable" + # class NetworkNICDisableAction(AbstractAction): # def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: # super().__init__(manager=manager) @@ -495,7 +497,7 @@ class ActionManager: "num_protocols": len(self.protocols), "num_ports": len(self.protocols), "num_ips": len(self.ip_address_list), - "max_acl_rules":max_acl_rules, + "max_acl_rules": max_acl_rules, "max_nics_per_node": max_nics_per_node, } self.actions: Dict[str, AbstractAction] = {} @@ -507,8 +509,8 @@ class ActionManager: # option_2: value2 # where `type` decides which AbstractAction subclass should be used # and `options` is an optional dict of options to pass to the init method of the action class - act_type = act_spec.get('type') - act_options = act_spec.get('options', {}) + act_type = act_spec.get("type") + act_options = act_spec.get("options", {}) self.actions[act_type] = self.__act_class_identifiers[act_type](self, **global_action_args, **act_options) self.action_map: Dict[int, Tuple[str, Dict]] = {} @@ -555,13 +557,12 @@ class ActionManager: param_combinations = list(itertools.product(*possibilities)) all_action_possibilities.extend( [ - ( - act_name, {param_names[i]:param_combinations[j][i] for i in range(len(param_names))} - ) for j in range(len(param_combinations))] - ) - - return {i:p for i,p in enumerate(all_action_possibilities)} + (act_name, {param_names[i]: param_combinations[j][i] for i in range(len(param_names))}) + for j in range(len(param_combinations)) + ] + ) + return {i: p for i, p in enumerate(all_action_possibilities)} def get_action(self, action: int) -> Tuple[str, Dict]: """Produce action in CAOS format""" @@ -627,7 +628,7 @@ class ActionManager: session=session, actions=cfg["action_list"], # node_uuids=cfg["options"]["node_uuids"], - **cfg['options'], + **cfg["options"], protocols=session.options.protocols, ports=session.options.ports, ip_address_list=None, diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 6083db6f..817e59b1 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -23,7 +23,7 @@ class AbstractAgent(ABC): observation_space: Optional[ObservationSpace], reward_function: Optional[RewardFunction], ) -> None: - self.agent_name:str = agent_name or "unnamed_agent" + self.agent_name: str = agent_name or "unnamed_agent" self.action_space: Optional[ActionManager] = action_space self.observation_space: Optional[ObservationSpace] = observation_space self.reward_function: Optional[RewardFunction] = reward_function @@ -46,9 +46,9 @@ class AbstractAgent(ABC): def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 0-39, # then use a bespoke conversion to take 1-40 int back into CAOS action - return ("DO_NOTHING", {} ) + return ("DO_NOTHING", {}) - def format_request(self, action:Tuple[str,Dict], options:Dict[str, int]) -> List[str]: + def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> List[str]: # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR """Format action into format expected by the simulator, and apply execution definition if applicable.""" diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 28c87af1..7b10f957 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -1,10 +1,11 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING, Sequence, Tuple +from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING from gym import spaces from pydantic import BaseModel from primaite.simulator.sim_container import Simulation + if TYPE_CHECKING: from primaite.game.session import PrimaiteSession @@ -29,7 +30,7 @@ def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: :return: The value in the dictionary :rtype: Any """ - key_list = [*keys] # copy keys to a new list to prevent editing original list + key_list = [*keys] # copy keys to a new list to prevent editing original list if len(key_list) == 0: return dictionary k = key_list.pop(0) @@ -58,7 +59,7 @@ class AbstractObservation(ABC): @classmethod @abstractmethod - def from_config(cls, config:Dict, session:"PrimaiteSession"): + def from_config(cls, config: Dict, session: "PrimaiteSession"): """Create this observation space component form a serialised format. The `session` parameter is for a the PrimaiteSession object that spawns this component. During deserialisation, @@ -98,7 +99,7 @@ class FileObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None): - return cls(where=parent_where+["files", config["file_name"]]) + return cls(where=parent_where + ["files", config["file_name"]]) class ServiceObservation(AbstractObservation): @@ -132,11 +133,8 @@ class ServiceObservation(AbstractObservation): return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]=None): - return cls( - where=parent_where+["services",session.ref_map_services[config['service_ref']].uuid] - ) - + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None): + return cls(where=parent_where + ["services", session.ref_map_services[config["service_ref"]].uuid]) class LinkObservation(AbstractObservation): @@ -179,7 +177,7 @@ class LinkObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession"): - return cls(where=['network','links', session.ref_map_links[config['link_ref']]]) + return cls(where=["network", "links", session.ref_map_links[config["link_ref"]]]) class FolderObservation(AbstractObservation): @@ -237,13 +235,13 @@ class FolderObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]): - where = parent_where + ["folders", config['folder_name']] + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]): + where = parent_where + ["folders", config["folder_name"]] file_configs = config["files"] files = [FileObservation.from_config(config=f, session=session, parent_where=where) for f in file_configs] - return cls(where=where,files=files) + return cls(where=where, files=files) class NicObservation(AbstractObservation): @@ -267,7 +265,7 @@ class NicObservation(AbstractObservation): return spaces.Dict({"nic_status": spaces.Discrete(3)}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]): return cls(where=parent_where + ["NICs", config["nic_uuid"]]) @@ -278,7 +276,7 @@ class NodeObservation(AbstractObservation): services: List[ServiceObservation] = [], folders: List[FolderObservation] = [], nics: List[NicObservation] = [], - logon_status:bool=False + logon_status: bool = False, ) -> None: """ Configurable observation for a node in the simulation. @@ -306,7 +304,7 @@ class NodeObservation(AbstractObservation): self.services: List[ServiceObservation] = services self.folders: List[FolderObservation] = folders self.nics: List[NicObservation] = nics - self.logon_status:bool=logon_status + self.logon_status: bool = logon_status self.default_observation: Dict = { "SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)}, @@ -315,7 +313,7 @@ class NodeObservation(AbstractObservation): "operating_status": 0, } if self.logon_status: - self.default_observation['logon_status']=0 + self.default_observation["logon_status"] = 0 def observe(self, state: Dict) -> Dict: if self.where is None: @@ -332,7 +330,7 @@ class NodeObservation(AbstractObservation): obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)} if self.logon_status: - obs['logon_status'] = 0 + obs["logon_status"] = 0 return obs @@ -345,26 +343,28 @@ class NodeObservation(AbstractObservation): "NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}), } if self.logon_status: - space_shape['logon_status'] = spaces.Discrete(3) + space_shape["logon_status"] = spaces.Discrete(3) return spaces.Dict(space_shape) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None) -> "NodeObservation": - node_uuid = session.ref_map_nodes[config['node_ref']] + def from_config( + cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None + ) -> "NodeObservation": + node_uuid = session.ref_map_nodes[config["node_ref"]] if parent_where is None: where = ["network", "nodes", node_uuid] else: where = parent_where + ["nodes", node_uuid] - svc_configs = config.get('services', {}) + svc_configs = config.get("services", {}) services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs] - folder_configs = config.get('folders', {}) - folders = [FolderObservation.from_config(config=c,session=session, parent_where=where) for c in folder_configs] + folder_configs = config.get("folders", {}) + folders = [FolderObservation.from_config(config=c, session=session, parent_where=where) for c in folder_configs] nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys() - nic_configs = [{'nic_uuid':n for n in nic_uuids }] if nic_uuids else [] + nic_configs = [{"nic_uuid": n for n in nic_uuids}] if nic_uuids else [] nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs] - logon_status = config.get('logon_status',False) + logon_status = config.get("logon_status", False) return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status) @@ -374,7 +374,12 @@ class AclObservation(AbstractObservation): # if a file is created at runtime, we have currently got no way of telling the observation space to track it. # this needs adding, but not for the MVP. def __init__( - self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[Tuple[str]] = None, num_rules: int = 10 + self, + node_ip_to_id: Dict[str, int], + ports: List[int], + protocols: list[str], + where: Optional[Tuple[str]] = None, + num_rules: int = 10, ) -> None: super().__init__() self.where: Optional[Tuple[str]] = where @@ -386,16 +391,18 @@ class AclObservation(AbstractObservation): self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)} "List of protocols which are part of the game, defines ordering when converting to an ID" self.default_observation: Dict = { - "RULES": {i+ 1:{ - "position": i, - "permission": 0, - "source_node_id": 0, - "source_port": 0, - "dest_node_id": 0, - "dest_port": 0, - "protocol": 0, + "RULES": { + i + + 1: { + "position": i, + "permission": 0, + "source_node_id": 0, + "source_port": 0, + "dest_node_id": 0, + "dest_port": 0, + "protocol": 0, } - for i in range(self.num_rules) + for i in range(self.num_rules) } } @@ -406,8 +413,7 @@ class AclObservation(AbstractObservation): if acl_state is NOT_PRESENT_IN_STATE: return self.default_observation - - #TODO: what if the ACL has more rules than num of max rules for obs space + # TODO: what if the ACL has more rules than num of max rules for obs space obs = {} obs["RULES"] = {} for i, rule_state in acl_state.items(): @@ -439,7 +445,8 @@ class AclObservation(AbstractObservation): { "RULE": spaces.Dict( { - i + 1: spaces.Dict( + i + + 1: spaces.Dict( { "position": spaces.Discrete(self.num_rules), "permission": spaces.Discrete(3), @@ -460,23 +467,23 @@ class AclObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation": node_ip_to_idx = {} - for node_idx, node_cfg in enumerate(config['node_order']): + for node_idx, node_cfg in enumerate(config["node_order"]): n_ref = node_cfg["node_ref"] n_obj = session.simulation.network.nodes[session.ref_map_nodes[n_ref]] for nic_uuid, nic_obj in n_obj.nics.items(): node_ip_to_idx[nic_obj.ip_address] = node_idx + 2 - router_uuid = session.ref_map_nodes[config['router_node_ref']] + router_uuid = session.ref_map_nodes[config["router_node_ref"]] return cls( node_ip_to_id=node_ip_to_idx, ports=session.options.ports, protocols=session.options.protocols, - where=["network", "nodes", router_uuid, "acl", "acl"]) - + where=["network", "nodes", router_uuid, "acl", "acl"], + ) class NullObservation(AbstractObservation): - def __init__(self, where:Optional[List[str]]=None): + def __init__(self, where: Optional[List[str]] = None): self.default_observation: Dict = {} def observe(self, state: Dict) -> Dict: @@ -487,20 +494,22 @@ class NullObservation(AbstractObservation): return spaces.Dict({}) @classmethod - def from_config(cls, config:Dict, session:Optional["PrimaiteSession"]=None) -> "NullObservation": + def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation": return cls() -class ICSObservation(NullObservation): pass + +class ICSObservation(NullObservation): + pass class UC2BlueObservation(AbstractObservation): def __init__( - self, - nodes: List[NodeObservation], - links: List[LinkObservation], - acl: AclObservation, - ics: ICSObservation, - where:Optional[List[str]] = None, + self, + nodes: List[NodeObservation], + links: List[LinkObservation], + acl: AclObservation, + ics: ICSObservation, + where: Optional[List[str]] = None, ) -> None: super().__init__() self.where: Optional[Tuple[str]] = where @@ -510,36 +519,38 @@ class UC2BlueObservation(AbstractObservation): self.acl: AclObservation = acl self.ics: ICSObservation = ics - self.default_observation : Dict = { - "NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)}, - "LINKS": {i+1: l.default_observation for i,l in enumerate(self.links)}, + self.default_observation: Dict = { + "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, + "LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)}, "ACL": self.acl.default_observation, "ICS": self.ics.default_observation, } - def observe(self, state:Dict) -> Dict: + def observe(self, state: Dict) -> Dict: if self.where is None: return self.default_observation obs = {} - obs['NODES'] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} - obs['LINKS'] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} - obs['ACL'] = self.acl.observe(state) - obs['ICS'] = self.ics.observe(state) + obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} + obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} + obs["ACL"] = self.acl.observe(state) + obs["ICS"] = self.ics.observe(state) return obs @property def space(self) -> spaces.Space: - return spaces.Dict({ - "NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}), - "LINKS": spaces.Dict({i+1: link.space for i, link in enumerate(self.links)}), - "ACL": self.acl.space, - "ICS": self.ics.space, - }) + return spaces.Dict( + { + "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), + "LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}), + "ACL": self.acl.space, + "ICS": self.ics.space, + } + ) @classmethod - def from_config(cls, config:Dict, session:"PrimaiteSession"): + def from_config(cls, config: Dict, session: "PrimaiteSession"): node_configs = config["nodes"] nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs] @@ -551,18 +562,18 @@ class UC2BlueObservation(AbstractObservation): ics_config = config["ics"] ics = ICSObservation.from_config(config=ics_config, session=session) - new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network']) + new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"]) return new class UC2RedObservation(AbstractObservation): - def __init__(self, nodes:List[NodeObservation], where:Optional[List[str]] = None) -> None: + def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None: super().__init__() - self.where:Optional[List[str]] = where + self.where: Optional[List[str]] = where self.nodes: List[NodeObservation] = nodes - self.default_observation : Dict = { - "NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)}, + self.default_observation: Dict = { + "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, } def observe(self, state: Dict) -> Dict: @@ -570,14 +581,16 @@ class UC2RedObservation(AbstractObservation): return self.default_observation obs = {} - obs['NODES'] = {i+1: node.observe(state) for i, node in enumerate(self.nodes)} + obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} return obs @property def space(self) -> spaces.Space: - return spaces.Dict({ - "NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}), - }) + return spaces.Dict( + { + "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), + } + ) @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession"): @@ -586,7 +599,9 @@ class UC2RedObservation(AbstractObservation): return cls(nodes=nodes, where=["network"]) -class UC2GreenObservation(NullObservation): pass +class UC2GreenObservation(NullObservation): + pass + class ObservationSpace: """ @@ -603,7 +618,7 @@ class ObservationSpace: # what this class does: # keep a list of observations # create observations for an actor from the config - def __init__(self, observation:AbstractObservation) -> None: + def __init__(self, observation: AbstractObservation) -> None: self.obs: AbstractObservation = observation def observe(self, state) -> Dict: @@ -614,12 +629,12 @@ class ObservationSpace: return self.obs.space @classmethod - def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace": - if config['type'] == "UC2BlueObservation": - return cls(UC2BlueObservation.from_config(config.get('options',{}), session=session)) - elif config['type'] == "UC2RedObservation": - return cls(UC2RedObservation.from_config(config.get('options',{}), session=session)) - elif config['type'] == "UC2GreenObservation": - return cls(UC2GreenObservation.from_config(config.get("options",{}), session=session)) + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "ObservationSpace": + if config["type"] == "UC2BlueObservation": + return cls(UC2BlueObservation.from_config(config.get("options", {}), session=session)) + elif config["type"] == "UC2RedObservation": + return cls(UC2RedObservation.from_config(config.get("options", {}), session=session)) + elif config["type"] == "UC2GreenObservation": + return cls(UC2GreenObservation.from_config(config.get("options", {}), session=session)) else: raise ValueError("Observation space type invalid") diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index a4ceb2dd..18925edc 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -1,34 +1,34 @@ from abc import ABC, abstractmethod from typing import Any, Dict, List -class AbstractReward(): + +class AbstractReward: def __init__(self): ... @abstractmethod - def calculate(self, state:Dict) -> float: + def calculate(self, state: Dict) -> float: return 0.3 -class DummyReward(AbstractReward): +class DummyReward(AbstractReward): def calculate(self, state: Dict) -> float: return -0.1 -class RewardFunction(): - __rew_class_identifiers:Dict[str,type[AbstractReward]] = { - "DUMMY" : DummyReward - } - def __init__(self, reward_function:AbstractReward): + +class RewardFunction: + __rew_class_identifiers: Dict[str, type[AbstractReward]] = {"DUMMY": DummyReward} + + def __init__(self, reward_function: AbstractReward): self.reward: AbstractReward = reward_function - def calculate(self, state:Dict) -> float: + def calculate(self, state: Dict) -> float: return self.reward.calculate(state) @classmethod - def from_config(cls, cfg:Dict) -> "RewardFunction": - for rew_component_cfg in cfg['reward_components']: - rew_type = rew_component_cfg['type'] + def from_config(cls, cfg: Dict) -> "RewardFunction": + for rew_component_cfg in cfg["reward_components"]: + rew_type = rew_component_cfg["type"] rew_component = cls.__rew_class_identifiers[rew_type]() new = cls(reward_function=rew_component) return new - diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 4bcf26e4..7b2225ef 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -10,6 +10,7 @@ from typing import Dict, List from pydantic import BaseModel +from primaite import getLogger from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent, RandomAgent from primaite.game.agent.observations import ( @@ -37,14 +38,12 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.sim_container import Simulation from primaite.simulator.system.applications.database_client import DatabaseClient -from primaite.simulator.system.services.database_service import DatabaseService -from primaite.simulator.system.services.dns_client import DNSClient -from primaite.simulator.system.services.dns_server import DNSServer +from primaite.simulator.system.services.database.database_service import DatabaseService +from primaite.simulator.system.services.dns.dns_client import DNSClient +from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.service import Service -from primaite import getLogger - _LOGGER = getLogger(__name__) @@ -91,7 +90,7 @@ class PrimaiteSession: agent_action, action_options = agent.get_action(agent_obs, agent_reward) # 9. CAOS action is converted into request (extra information might be needed to enrich # the request, this is what the execution definition is there for) - _LOGGER.debug(f"Formatting agent action {agent_action}") # maybe too many debug log statements + _LOGGER.debug(f"Formatting agent action {agent_action}") # maybe too many debug log statements agent_request = agent.format_request(agent_action, action_options) # 10. primaite session receives the action from the agents and asks the simulation to apply each @@ -106,8 +105,8 @@ class PrimaiteSession: def from_config(cls, cfg: dict) -> "PrimaiteSession": sess = cls() sess.options = PrimaiteSessionOptions( - ports = cfg['game_config']['ports'], - protocols = cfg['game_config']['protocols'], + ports=cfg["game_config"]["ports"], + protocols=cfg["game_config"]["protocols"], ) sim = sess.simulation net = sim.network @@ -230,7 +229,7 @@ class PrimaiteSession: reward_function_cfg = agent_cfg["reward_function"] # CREATE OBSERVATION SPACE - obs_space=ObservationSpace.from_config(observation_space_cfg, sess) + obs_space = ObservationSpace.from_config(observation_space_cfg, sess) """ # if observation_space_cfg is None: @@ -331,23 +330,23 @@ class PrimaiteSession: """ # CREATE ACTION SPACE - action_space_cfg['options']['node_uuids'] = [] + action_space_cfg["options"]["node_uuids"] = [] # if a list of nodes is defined, convert them from node references to node UUIDs - for action_node_option in action_space_cfg.get('options',{}).pop('nodes', {}): - if 'node_ref' in action_node_option: - node_uuid = sess.ref_map_nodes[action_node_option['node_ref']] - action_space_cfg['options']['node_uuids'].append(node_uuid) + for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}): + if "node_ref" in action_node_option: + node_uuid = sess.ref_map_nodes[action_node_option["node_ref"]] + action_space_cfg["options"]["node_uuids"].append(node_uuid) # Each action space can potentially have a different list of nodes that it can apply to. Therefore, # we will pass node_uuids as a part of the action space config. # However, it's not possible to specify the node uuids directly in the config, as they are generated # dynamically, so we have to translate node references to uuids before passing this config on. - if 'action_list' in action_space_cfg: - for action_config in action_space_cfg['action_list']: - if 'options' in action_config: - if 'target_router_ref' in action_config['options']: - _target = action_config['options']['target_router_ref'] - action_config['options']['target_router_uuid'] = sess.ref_map_nodes[_target] + if "action_list" in action_space_cfg: + for action_config in action_space_cfg["action_list"]: + if "options" in action_config: + if "target_router_ref" in action_config["options"]: + _target = action_config["options"]["target_router_ref"] + action_config["options"]["target_router_uuid"] = sess.ref_map_nodes[_target] action_space = ActionManager.from_config(sess, action_space_cfg) @@ -357,16 +356,30 @@ class PrimaiteSession: # CREATE AGENT if agent_type == "GreenWebBrowsingAgent": # TODO: implement non-random agents and fix this parsing - new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) + new_agent = RandomAgent( + agent_name=agent_cfg["ref"], + action_space=action_space, + observation_space=obs_space, + reward_function=rew_function, + ) sess.agents.append(new_agent) elif agent_type == "GATERLAgent": - new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) + new_agent = RandomAgent( + agent_name=agent_cfg["ref"], + action_space=action_space, + observation_space=obs_space, + reward_function=rew_function, + ) sess.agents.append(new_agent) elif agent_type == "RedDatabaseCorruptingAgent": - new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function) + new_agent = RandomAgent( + agent_name=agent_cfg["ref"], + action_space=action_space, + observation_space=obs_space, + reward_function=rew_function, + ) sess.agents.append(new_agent) else: print("agent type not found") - return sess diff --git a/src/primaite/simulator/network/hardware/nodes/switch.py b/src/primaite/simulator/network/hardware/nodes/switch.py index bb296203..09b53483 100644 --- a/src/primaite/simulator/network/hardware/nodes/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/switch.py @@ -57,8 +57,8 @@ class Switch(Node): """ state = super().describe_state() state["ports"] = {port_num: port.describe_state() for port_num, port in self.switch_ports.items()} - state["num_ports"]= self.num_ports # redundant? - state["mac_address_table"]= {mac: port for mac, port in self.mac_address_table.items()} + state["num_ports"] = self.num_ports # redundant? + state["mac_address_table"] = {mac: port for mac, port in self.mac_address_table.items()} return state def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort): From b53c3856dd470e88fb18d8635758d9374398dfea Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 10 Oct 2023 09:48:04 +0100 Subject: [PATCH 19/35] Add GATE wheel temporarily --- example_config.yaml | 8 +- src/primaite/game/agent/GATE_agents.py | 59 ++++++++ src/primaite/game/session.py | 178 +++++++++++------------- src/primaite/utils/start_gate_server.py | 5 + 4 files changed, 149 insertions(+), 101 deletions(-) create mode 100644 src/primaite/utils/start_gate_server.py diff --git a/example_config.yaml b/example_config.yaml index f7faf589..afdf1b0a 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -1,8 +1,12 @@ training_config: rl_framework: SB3 - rl_algo: PPO + rl_algorithm: PPO + seed: 333 + n_learn_episodes: 2 n_learn_steps: 128 - n_learn_episodes: 1000 + n_eval_episodes: 2 + n_eval_steps: 128 + game_config: ports: diff --git a/src/primaite/game/agent/GATE_agents.py b/src/primaite/game/agent/GATE_agents.py index 5bdfebe4..ac1d776b 100644 --- a/src/primaite/game/agent/GATE_agents.py +++ b/src/primaite/game/agent/GATE_agents.py @@ -1,5 +1,64 @@ from primaite.game.agent.interface import AbstractGATEAgent +from arcd_gate.client.gate_client import GATEClient +class GATEMan(GATEClient): + + @property + def rl_framework(self) -> str: + return "SB3" + + @property + def rl_framework(self) -> str: + pass + + @property + def rl_algorithm(self) -> str: + pass + + @property + def seed(self) -> Optional[int]: + return None + + @property + def n_learn_episodes(self) -> int: + return 0 + + @property + def n_learn_steps(self) -> int: + return 0 + + @property + def n_eval_episodes(self) -> int: + return 0 + + @property + def n_eval_steps(self) -> int: + return 0 + + @property + def action_space(self) -> spaces.Space: + pass + + @property + def observation_space(self) -> spaces.Space: + pass + + def step(self, action: ActType) -> Tuple[np.ndarray, float, bool, bool, Dict]: + pass + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None) -> Tuple[np.ndarray, Dict]: + pass + + def close(self): + pass + class GATERLAgent(AbstractGATEAgent): ... + # The communication with GATE needs to be handled by the PrimaiteSession, rather than by individual agents, + # because when we are supporting MARL, the actions form multiple agents will have to be batched + + # For example MultiAgentEnv in Ray allows sending a dict of observations of multiple agents, then it will reply + # with the actions for those agents. + + diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 7b2225ef..7746f78c 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -6,7 +6,9 @@ # 5. idk from ipaddress import IPv4Address -from typing import Dict, List +from typing import Any, Dict, List, Optional, Tuple +from gymnasium.vector.utils import spaces +import numpy as np from pydantic import BaseModel @@ -44,26 +46,101 @@ from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.service import Service +from arcd_gate.client.gate_client import GATEClient, ActType +from numpy import ndarray + _LOGGER = getLogger(__name__) +class PrimaiteGATEClient(GATEClient): + def __init__(self, parent_session:"PrimaiteSession", service_port: int = 50000): + super().__init__(service_port=service_port) + self.parent_session:"PrimaiteSession" + + @property + def rl_framework(self) -> str: + return self.parent_session.training_options.rl_framework + + @property + def rl_algorithm(self) -> str: + return self.parent_session.training_options.rl_algorithm + + @property + def seed(self) -> int | None: + return self.parent_session.training_options.seed + + @property + def n_learn_episodes(self) -> int: + return self.parent_session.training_options.n_learn_episodes + + @property + def n_learn_steps(self) -> int: + return self.parent_session.training_options.n_learn_steps + + @property + def n_eval_episodes(self) -> int: + return self.parent_session.training_options.n_eval_episodes + + @property + def n_eval_steps(self) -> int: + return self.parent_session.training_options.n_eval_steps + + @property + def action_space(self) -> spaces.Space: + return self.parent_session.rl_agent.action_space + + @property + def observation_space(self) -> spaces.Space: + return self.parent_session.rl_agent.observation_space + + def step(self, action: ActType) -> Tuple[ndarray, float, bool, bool, Dict]: + self.parent_session.step() + #TODO: not sure how to go about this. + + def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ndarray, Dict]: + ... + + def close(self): + ... class PrimaiteSessionOptions(BaseModel): ports: List[str] protocols: List[str] +class TrainingOptions(BaseModel): + rl_framework:str + rl_algorithm:str + seed:Optional[int] + n_learn_episodes:int + n_learn_steps:int + n_eval_episodes:int + n_eval_steps:int + + class PrimaiteSession: def __init__(self): self.simulation: Simulation = Simulation() self.agents: List[AbstractAgent] = [] + self.rl_agent: AbstractAgent + # which of the agents should be used for sending RL data to GATE client? self.step_counter: int = 0 self.episode_counter: int = 0 self.options: PrimaiteSessionOptions + self.training_options: TrainingOptions self.ref_map_nodes: Dict[str, Node] = {} self.ref_map_services: Dict[str, Service] = {} self.ref_map_links: Dict[str, Link] = {} + def start_session(self, opts="TODO..."): + """Commence the session, this gives the gate client control over the simulation/agent loop.""" + ... + + def eval(self, opts="TODO..."): + ... + + + def step(self): _LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}") # currently designed with assumption that all agents act once per step in order @@ -108,6 +185,7 @@ class PrimaiteSession: ports=cfg["game_config"]["ports"], protocols=cfg["game_config"]["protocols"], ) + sess.training_options = TrainingOptions(**cfg['training_config']) sim = sess.simulation net = sim.network @@ -231,104 +309,6 @@ class PrimaiteSession: # CREATE OBSERVATION SPACE obs_space = ObservationSpace.from_config(observation_space_cfg, sess) - """ - # if observation_space_cfg is None: - # obs_space = NullObservation() - # elif observation_space_cfg["type"] == "UC2BlueObservation": - # node_obs_list = [] - # link_obs_list = [] - - # # node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses - # node_ip_to_index = {} - # for node_idx, node_cfg in enumerate(nodes_cfg): - # n_ref = node_cfg["ref"] - # n_obj = net.nodes[ref_map_nodes[n_ref]] - # for nic_uuid, nic_obj in n_obj.nics.items(): - # node_ip_to_index[nic_obj.ip_address] = node_idx + 2 - - # for node_obs_cfg in observation_space_cfg["options"]["nodes"]: - # node_ref = node_obs_cfg["node_ref"] - # folder_obs_list = [] - # service_obs_list = [] - # if "services" in node_obs_cfg: - # for service_obs_cfg in node_obs_cfg["services"]: - # service_obs_list.append( - # ServiceObservation( - # where=[ - # "network", - # "nodes", - # ref_map_nodes[node_ref], - # "services", - # ref_map_services[service_obs_cfg["service_ref"]], - # ] - # ) - # ) - # if "folders" in node_obs_cfg: - # for folder_obs_cfg in node_obs_cfg["folders"]: - # file_obs_list = [] - # if "files" in folder_obs_cfg: - # for file_obs_cfg in folder_obs_cfg["files"]: - # file_obs_list.append( - # FileObservation( - # where=[ - # "network", - # "nodes", - # ref_map_nodes[node_ref], - # "folders", - # folder_obs_cfg["folder_name"], - # "files", - # file_obs_cfg["file_name"], - # ] - # ) - # ) - # folder_obs_list.append( - # FolderObservation( - # where=[ - # "network", - # "nodes", - # ref_map_nodes[node_ref], - # "folders", - # folder_obs_cfg["folder_name"], - # ], - # files=file_obs_list, - # ) - # ) - # nic_obs_list = [] - # for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys(): - # nic_obs_list.append( - # NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid]) - # ) - # node_obs_list.append( - # NodeObservation( - # where=["network", "nodes", ref_map_nodes[node_ref]], - # services=service_obs_list, - # folders=folder_obs_list, - # nics=nic_obs_list, - # logon_status=False, - # ) - # ) - # for link_obs_cfg in observation_space_cfg["options"]["links"]: - # link_ref = link_obs_cfg["link_ref"] - # link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]])) - - # acl_obs = AclObservation( - # node_ip_to_id=node_ip_to_index, - # ports=game_cfg["ports"], - # protocols=game_cfg["ports"], - # where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]], - # ) - # obs_space = UC2BlueObservation( - # nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation() - # ) - # elif observation_space_cfg["type"] == "UC2RedObservation": - # obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim) - # elif observation_space_cfg["type"] == "UC2GreenObservation": - # obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{})) - # else: - # print("observation space config not specified correctly.") - # obs_space = NullObservation() - """ - # CREATE ACTION SPACE action_space_cfg["options"]["node_uuids"] = [] # if a list of nodes is defined, convert them from node references to node UUIDs diff --git a/src/primaite/utils/start_gate_server.py b/src/primaite/utils/start_gate_server.py new file mode 100644 index 00000000..53508cd2 --- /dev/null +++ b/src/primaite/utils/start_gate_server.py @@ -0,0 +1,5 @@ +"""Utility script to start the gate server for running PrimAITE in attached mode.""" +from arcd_gate.server.gate_service import GATEService + +service = GATEService() +service.start() From def2c3699be661656682748677fbcb7a34de5816 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 10 Oct 2023 21:01:09 +0100 Subject: [PATCH 20/35] Add empty obs placeholders --- example_config.yaml | 22 +- sandbox.ipynb | 822 ++++++++++++++++++++++- sandbox.py | 19 + src/primaite/environment/observations.py | 2 +- src/primaite/game/agent/GATE_agents.py | 65 +- src/primaite/game/agent/actions.py | 2 +- src/primaite/game/agent/observations.py | 82 ++- src/primaite/game/session.py | 49 +- src/primaite/transactions/transaction.py | 2 +- 9 files changed, 964 insertions(+), 101 deletions(-) create mode 100644 sandbox.py diff --git a/example_config.yaml b/example_config.yaml index afdf1b0a..a35c82e0 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -12,10 +12,12 @@ game_config: ports: - ARP - DNS + - HTTP - POSTGRES_SERVER protocols: - ICMP - TCP + - UDP agents: - ref: client_1_green_user @@ -111,10 +113,10 @@ game_config: observation_space: type: UC2BlueObservation options: + num_services_per_node: 2 + num_folders_per_node: 2 + num_files_per_folder: 2 nodes: - - node_ref: router_1 #TODO: more sub-options here - - node_ref: switch_1 - - node_ref: switch_2 - node_ref: domain_controller services: - service_ref: domain_controller_dns_server @@ -147,17 +149,23 @@ game_config: - link_ref: switch_2___security_suite acl: router_node_ref: router_1 - node_order: - - node_ref: router_1 - - node_ref: switch_1 - - node_ref: switch_2 + ip_address_order: - node_ref: domain_controller + nic_num: 1 - node_ref: web_server + nic_num: 1 - node_ref: database_server + nic_num: 1 - node_ref: backup_server + nic_num: 1 - node_ref: security_suite + nic_num: 1 - node_ref: client_1 + nic_num: 1 - node_ref: client_2 + nic_num: 1 + - node_ref: security_suite + nic_num: 2 ics: null action_space: diff --git a/sandbox.ipynb b/sandbox.ipynb index a2150921..73c3e682 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -2,9 +2,105 @@ "cells": [ { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-10 21:00:40,310: Added node f0fb4743-43d5-4741-a0e5-78340a458a11 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,313: Added node 86f4ee9d-386d-4436-96e7-d00dd8bae746 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,316: Added node d2b38f6a-10fe-4d28-86a5-a1d2010c4bea to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,319: Added node 29d55bd4-d59e-4f73-95ad-b430c8716b15 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,323: Added node 5d266f99-03a6-47b3-ab0b-28b8a6813a11 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,333: Added node e4f397c7-eed2-4e69-afaf-ce44664ff1d2 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,338: Added node 19dbb2d4-648b-4387-aa15-99757b513acb to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,344: Added node f508ea65-e660-4b91-8628-5c586075137b to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,350: Added node ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,356: Added node 0ce9efc6-39ae-4d3d-82ad-43be5ac57586 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", + "2023-10-10 21:00:40,360: NIC ba:25:b9:f4:b0:74/192.168.1.1 connected to Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36\n", + "2023-10-10 21:00:40,361: SwitchPort 3a:25:73:6c:4c:36 connected to Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36\n", + "2023-10-10 21:00:40,363: Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36 up\n", + "2023-10-10 21:00:40,364: Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36 up\n", + "2023-10-10 21:00:40,366: Added link eda2befd-9f68-440c-9b20-06838d9e6553 to connect ba:25:b9:f4:b0:74/192.168.1.1 and 3a:25:73:6c:4c:36\n", + "2023-10-10 21:00:40,367: NIC c0:ce:5a:6d:71:73/192.168.1.1 connected to Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3\n", + "2023-10-10 21:00:40,368: SwitchPort b2:a4:b3:fc:da:c3 connected to Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3\n", + "2023-10-10 21:00:40,370: Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3 up\n", + "2023-10-10 21:00:40,370: Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3 up\n", + "2023-10-10 21:00:40,372: Added link 0e2a1df6-1c3c-4517-b41f-04a5cefe307c to connect c0:ce:5a:6d:71:73/192.168.1.1 and b2:a4:b3:fc:da:c3\n", + "2023-10-10 21:00:40,373: SwitchPort 72:f3:93:a5:a5:59 connected to Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10\n", + "2023-10-10 21:00:40,377: Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10 up\n", + "2023-10-10 21:00:40,379: NIC e1:9b:c0:59:1d:46/192.168.1.10 connected to Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10\n", + "2023-10-10 21:00:40,380: Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10 up\n", + "2023-10-10 21:00:40,381: Added link 48c9173d-847b-441f-9f15-f7838cf09083 to connect 72:f3:93:a5:a5:59 and e1:9b:c0:59:1d:46/192.168.1.10\n", + "2023-10-10 21:00:40,382: SwitchPort 67:36:e8:51:35:f2 connected to Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12\n", + "2023-10-10 21:00:40,384: Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12 up\n", + "2023-10-10 21:00:40,385: NIC 3b:a3:b4:ec:a0:2a/192.168.1.12 connected to Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12\n", + "2023-10-10 21:00:40,386: Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12 up\n", + "2023-10-10 21:00:40,386: Added link 7259c2a4-44f7-44d9-87de-1a692c98ac58 to connect 67:36:e8:51:35:f2 and 3b:a3:b4:ec:a0:2a/192.168.1.12\n", + "2023-10-10 21:00:40,388: SwitchPort be:56:a8:d3:f9:6c connected to Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14\n", + "2023-10-10 21:00:40,391: Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14 up\n", + "2023-10-10 21:00:40,392: NIC 99:24:3a:ad:99:5b/192.168.1.14 connected to Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14\n", + "2023-10-10 21:00:40,394: Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14 up\n", + "2023-10-10 21:00:40,396: Added link ff3c3f9a-c267-421a-9e2f-bcc11a6fa082 to connect be:56:a8:d3:f9:6c and 99:24:3a:ad:99:5b/192.168.1.14\n", + "2023-10-10 21:00:40,397: SwitchPort 4f:91:78:5a:68:3f connected to Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16\n", + "2023-10-10 21:00:40,399: Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16 up\n", + "2023-10-10 21:00:40,400: NIC 2b:ff:81:44:7f:61/192.168.1.16 connected to Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16\n", + "2023-10-10 21:00:40,401: Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16 up\n", + "2023-10-10 21:00:40,402: Added link b5a2e61c-f14d-45a3-a27e-2451f21229d5 to connect 4f:91:78:5a:68:3f and 2b:ff:81:44:7f:61/192.168.1.16\n", + "2023-10-10 21:00:40,403: SwitchPort 14:8e:63:89:28:f6 connected to Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110\n", + "2023-10-10 21:00:40,404: Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110 up\n", + "2023-10-10 21:00:40,405: NIC 72:ca:83:f9:d0:66/192.168.1.110 connected to Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110\n", + "2023-10-10 21:00:40,406: Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110 up\n", + "2023-10-10 21:00:40,407: Added link d0906ccd-5f98-415a-b08e-b8ebfdd3c5a9 to connect 14:8e:63:89:28:f6 and 72:ca:83:f9:d0:66/192.168.1.110\n", + "2023-10-10 21:00:40,409: SwitchPort e5:20:55:91:b5:b9 connected to Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21\n", + "2023-10-10 21:00:40,412: Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21 up\n", + "2023-10-10 21:00:40,413: NIC 27:2f:d1:75:04:f9/192.168.10.21 connected to Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21\n", + "2023-10-10 21:00:40,414: Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21 up\n", + "2023-10-10 21:00:40,415: Added link 51f14a80-8fd6-4de9-86a3-9c726c036167 to connect e5:20:55:91:b5:b9 and 27:2f:d1:75:04:f9/192.168.10.21\n", + "2023-10-10 21:00:40,416: SwitchPort 4e:58:c4:9c:0c:6d connected to Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22\n", + "2023-10-10 21:00:40,418: Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22 up\n", + "2023-10-10 21:00:40,420: NIC 2e:f5:71:bb:73:ec/192.168.10.22 connected to Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22\n", + "2023-10-10 21:00:40,420: Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22 up\n", + "2023-10-10 21:00:40,421: Added link 7463a72f-af42-435f-ad98-83bcfe5728a3 to connect 4e:58:c4:9c:0c:6d and 2e:f5:71:bb:73:ec/192.168.10.22\n", + "2023-10-10 21:00:40,423: SwitchPort 66:98:bd:ca:08:b0 connected to Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110\n", + "2023-10-10 21:00:40,425: Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110 up\n", + "2023-10-10 21:00:40,427: NIC cf:f7:fc:d1:f4:ae/192.168.10.110 connected to Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110\n", + "2023-10-10 21:00:40,428: Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110 up\n", + "2023-10-10 21:00:40,430: Added link ff6fd52d-f893-4c0d-99d2-28f10c128043 to connect 66:98:bd:ca:08:b0 and cf:f7:fc:d1:f4:ae/192.168.10.110\n", + "2023-10-10 21:00:40,431: Stepping primaite session. Step counter: 0\n", + "2023-10-10 21:00:40,432: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,433: Getting agent action\n", + "2023-10-10 21:00:40,435: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,436: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,437: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,439: Getting agent action\n", + "2023-10-10 21:00:40,440: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,441: Sending request to simulation: ['do_nothing']\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "/home/cade/.local/state/primaite/2.0.0/log\n", + "service type not found DatabaseBackup\n", + "service type not found WebBrowser\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-10 21:00:40,443: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,445: Getting agent action\n", + "2023-10-10 21:00:40,447: Formatting agent action NODE_FOLDER_RESTORE\n", + "2023-10-10 21:00:40,449: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,450: Initiating simulation step 0\n" + ] + } + ], "source": [ "%load_ext autoreload\n", "%autoreload 2\n", @@ -14,11 +110,7 @@ "import logging\n", "_PRIMAITE_CONFIG['log_level']=logging.DEBUG\n", "print(PRIMAITE_PATHS.app_log_dir_path)\n", - "import itertools\n", "from primaite.game.session import PrimaiteSession\n", - "from primaite.simulator.sim_container import Simulation\n", - "from primaite.game.agent.interface import AbstractAgent\n", - "from primaite.simulator.network.networks import arcd_uc2_network\n", "import yaml\n", "\n", "with open('example_config.yaml', 'r') as file:\n", @@ -29,13 +121,727 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-10-10 21:00:40,481: Stepping primaite session. Step counter: 1\n", + "2023-10-10 21:00:40,484: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,486: Getting agent action\n", + "2023-10-10 21:00:40,487: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,488: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,489: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,491: Getting agent action\n", + "2023-10-10 21:00:40,493: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,494: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,495: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,497: Getting agent action\n", + "2023-10-10 21:00:40,498: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:40,499: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 8]\n", + "2023-10-10 21:00:40,500: Initiating simulation step 1\n", + "2023-10-10 21:00:40,502: Stepping primaite session. Step counter: 2\n", + "2023-10-10 21:00:40,503: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,504: Getting agent action\n", + "2023-10-10 21:00:40,505: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,505: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,506: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,508: Getting agent action\n", + "2023-10-10 21:00:40,511: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,513: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,513: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,515: Getting agent action\n", + "2023-10-10 21:00:40,516: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:40,518: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,519: Initiating simulation step 2\n", + "2023-10-10 21:00:40,520: Stepping primaite session. Step counter: 3\n", + "2023-10-10 21:00:40,521: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,522: Getting agent action\n", + "2023-10-10 21:00:40,524: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,525: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,526: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,528: Getting agent action\n", + "2023-10-10 21:00:40,530: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,531: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,533: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,534: Getting agent action\n", + "2023-10-10 21:00:40,535: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:40,537: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 1]\n", + "2023-10-10 21:00:40,538: Initiating simulation step 3\n", + "2023-10-10 21:00:40,539: Stepping primaite session. Step counter: 4\n", + "2023-10-10 21:00:40,541: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,543: Getting agent action\n", + "2023-10-10 21:00:40,545: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,546: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,547: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,550: Getting agent action\n", + "2023-10-10 21:00:40,552: Formatting agent action NODE_OS_SCAN\n", + "2023-10-10 21:00:40,554: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", + "2023-10-10 21:00:40,555: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,573: Getting agent action\n", + "2023-10-10 21:00:40,575: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:40,577: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 2]\n", + "2023-10-10 21:00:40,578: Initiating simulation step 4\n", + "2023-10-10 21:00:40,580: Stepping primaite session. Step counter: 5\n", + "2023-10-10 21:00:40,581: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,583: Getting agent action\n", + "2023-10-10 21:00:40,585: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,586: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,587: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,597: Getting agent action\n", + "2023-10-10 21:00:40,598: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,599: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,601: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,603: Getting agent action\n", + "2023-10-10 21:00:40,605: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:40,606: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 0]\n", + "2023-10-10 21:00:40,613: Initiating simulation step 5\n", + "2023-10-10 21:00:40,614: Stepping primaite session. Step counter: 6\n", + "2023-10-10 21:00:40,615: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,617: Getting agent action\n", + "2023-10-10 21:00:40,618: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,619: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,620: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,621: Getting agent action\n", + "2023-10-10 21:00:40,623: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,624: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,625: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,626: Getting agent action\n", + "2023-10-10 21:00:40,628: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:40,629: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 7]\n", + "2023-10-10 21:00:40,630: Initiating simulation step 6\n", + "2023-10-10 21:00:40,631: Stepping primaite session. Step counter: 7\n", + "2023-10-10 21:00:40,631: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,633: Getting agent action\n", + "2023-10-10 21:00:40,634: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,635: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,635: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,636: Getting agent action\n", + "2023-10-10 21:00:40,639: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,640: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,642: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,644: Getting agent action\n", + "2023-10-10 21:00:40,645: Formatting agent action NODE_SERVICE_STOP\n", + "2023-10-10 21:00:40,646: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,647: Initiating simulation step 7\n", + "2023-10-10 21:00:40,648: Stepping primaite session. Step counter: 8\n", + "2023-10-10 21:00:40,649: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,651: Getting agent action\n", + "2023-10-10 21:00:40,652: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,652: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,653: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,655: Getting agent action\n", + "2023-10-10 21:00:40,656: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,657: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,659: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,661: Getting agent action\n", + "2023-10-10 21:00:40,662: Formatting agent action NODE_FILE_RESTORE\n", + "2023-10-10 21:00:40,663: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,664: Initiating simulation step 8\n", + "2023-10-10 21:00:40,665: Stepping primaite session. Step counter: 9\n", + "2023-10-10 21:00:40,667: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,668: Getting agent action\n", + "2023-10-10 21:00:40,669: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,670: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,672: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,674: Getting agent action\n", + "2023-10-10 21:00:40,676: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,677: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,678: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,679: Getting agent action\n", + "2023-10-10 21:00:40,681: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,682: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,683: Initiating simulation step 9\n", + "2023-10-10 21:00:40,684: Stepping primaite session. Step counter: 10\n", + "2023-10-10 21:00:40,685: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,687: Getting agent action\n", + "2023-10-10 21:00:40,688: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,689: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,690: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,694: Getting agent action\n", + "2023-10-10 21:00:40,696: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,697: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,698: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,700: Getting agent action\n", + "2023-10-10 21:00:40,702: Formatting agent action NODE_FILE_SCAN\n", + "2023-10-10 21:00:40,705: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,706: Initiating simulation step 10\n", + "2023-10-10 21:00:40,709: Stepping primaite session. Step counter: 11\n", + "2023-10-10 21:00:40,711: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,713: Getting agent action\n", + "2023-10-10 21:00:40,715: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,716: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,717: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,719: Getting agent action\n", + "2023-10-10 21:00:40,722: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,724: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,726: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,728: Getting agent action\n", + "2023-10-10 21:00:40,730: Formatting agent action NODE_SERVICE_START\n", + "2023-10-10 21:00:40,731: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,733: Initiating simulation step 11\n", + "2023-10-10 21:00:40,735: Stepping primaite session. Step counter: 12\n", + "2023-10-10 21:00:40,736: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,738: Getting agent action\n", + "2023-10-10 21:00:40,739: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,740: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,741: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,743: Getting agent action\n", + "2023-10-10 21:00:40,746: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,748: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,749: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,752: Getting agent action\n", + "2023-10-10 21:00:40,755: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:40,758: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:40,760: Initiating simulation step 12\n", + "2023-10-10 21:00:40,762: Stepping primaite session. Step counter: 13\n", + "2023-10-10 21:00:40,763: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,766: Getting agent action\n", + "2023-10-10 21:00:40,768: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,769: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,771: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,773: Getting agent action\n", + "2023-10-10 21:00:40,777: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,779: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,780: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,782: Getting agent action\n", + "2023-10-10 21:00:40,784: Formatting agent action NODE_SERVICE_RESTART\n", + "2023-10-10 21:00:40,785: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,786: Initiating simulation step 13\n", + "2023-10-10 21:00:40,787: Stepping primaite session. Step counter: 14\n", + "2023-10-10 21:00:40,788: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,790: Getting agent action\n", + "2023-10-10 21:00:40,792: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,794: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,795: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,797: Getting agent action\n", + "2023-10-10 21:00:40,799: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,800: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,801: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,803: Getting agent action\n", + "2023-10-10 21:00:40,805: Formatting agent action NODE_SERVICE_DISABLE\n", + "2023-10-10 21:00:40,806: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,807: Initiating simulation step 14\n", + "2023-10-10 21:00:40,808: Stepping primaite session. Step counter: 15\n", + "2023-10-10 21:00:40,809: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,813: Getting agent action\n", + "2023-10-10 21:00:40,815: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,817: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,818: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,821: Getting agent action\n", + "2023-10-10 21:00:40,822: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,824: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,828: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,830: Getting agent action\n", + "2023-10-10 21:00:40,832: Formatting agent action NETWORK_NIC_DISABLE\n", + "2023-10-10 21:00:40,833: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,835: Initiating simulation step 15\n", + "2023-10-10 21:00:40,836: Stepping primaite session. Step counter: 16\n", + "2023-10-10 21:00:40,838: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,840: Getting agent action\n", + "2023-10-10 21:00:40,848: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,852: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,856: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,858: Getting agent action\n", + "2023-10-10 21:00:40,863: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,868: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,869: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,872: Getting agent action\n", + "2023-10-10 21:00:40,878: Formatting agent action NODE_OS_SCAN\n", + "2023-10-10 21:00:40,883: Sending request to simulation: ['network', 'node', '29d55bd4-d59e-4f73-95ad-b430c8716b15', 'scan']\n", + "2023-10-10 21:00:40,885: Initiating simulation step 16\n", + "2023-10-10 21:00:40,887: Stepping primaite session. Step counter: 17\n", + "2023-10-10 21:00:40,889: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,900: Getting agent action\n", + "2023-10-10 21:00:40,904: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,909: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,911: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,913: Getting agent action\n", + "2023-10-10 21:00:40,915: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,917: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,918: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,920: Getting agent action\n", + "2023-10-10 21:00:40,921: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:40,922: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:40,924: Initiating simulation step 17\n", + "2023-10-10 21:00:40,925: Stepping primaite session. Step counter: 18\n", + "2023-10-10 21:00:40,927: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,929: Getting agent action\n", + "2023-10-10 21:00:40,931: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,933: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,934: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,936: Getting agent action\n", + "2023-10-10 21:00:40,938: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,940: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,941: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,945: Getting agent action\n", + "2023-10-10 21:00:40,947: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:40,948: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,949: Initiating simulation step 18\n", + "2023-10-10 21:00:40,951: Stepping primaite session. Step counter: 19\n", + "2023-10-10 21:00:40,952: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,954: Getting agent action\n", + "2023-10-10 21:00:40,955: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,957: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,960: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,962: Getting agent action\n", + "2023-10-10 21:00:40,964: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:40,966: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,967: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,969: Getting agent action\n", + "2023-10-10 21:00:40,971: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,972: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,973: Initiating simulation step 19\n", + "2023-10-10 21:00:40,975: Stepping primaite session. Step counter: 20\n", + "2023-10-10 21:00:40,976: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:40,979: Getting agent action\n", + "2023-10-10 21:00:40,981: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:40,982: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,983: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:40,986: Getting agent action\n", + "2023-10-10 21:00:40,987: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:40,989: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:40,990: Sending simulation state to agent defender\n", + "2023-10-10 21:00:40,992: Getting agent action\n", + "2023-10-10 21:00:40,994: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:40,996: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:40,997: Initiating simulation step 20\n", + "2023-10-10 21:00:40,998: Stepping primaite session. Step counter: 21\n", + "2023-10-10 21:00:41,000: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,001: Getting agent action\n", + "2023-10-10 21:00:41,003: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,004: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,005: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,006: Getting agent action\n", + "2023-10-10 21:00:41,009: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,010: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,011: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,014: Getting agent action\n", + "2023-10-10 21:00:41,016: Formatting agent action NODE_FILE_REPAIR\n", + "2023-10-10 21:00:41,017: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,018: Initiating simulation step 21\n", + "2023-10-10 21:00:41,020: Stepping primaite session. Step counter: 22\n", + "2023-10-10 21:00:41,021: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,023: Getting agent action\n", + "2023-10-10 21:00:41,024: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,025: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,027: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,030: Getting agent action\n", + "2023-10-10 21:00:41,031: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,033: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,034: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,036: Getting agent action\n", + "2023-10-10 21:00:41,037: Formatting agent action NODE_FILE_REPAIR\n", + "2023-10-10 21:00:41,038: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,039: Initiating simulation step 22\n", + "2023-10-10 21:00:41,040: Stepping primaite session. Step counter: 23\n", + "2023-10-10 21:00:41,042: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,045: Getting agent action\n", + "2023-10-10 21:00:41,047: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,049: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,050: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,056: Getting agent action\n", + "2023-10-10 21:00:41,065: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,067: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,068: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,072: Getting agent action\n", + "2023-10-10 21:00:41,073: Formatting agent action NODE_FOLDER_REPAIR\n", + "2023-10-10 21:00:41,075: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,077: Initiating simulation step 23\n", + "2023-10-10 21:00:41,079: Stepping primaite session. Step counter: 24\n", + "2023-10-10 21:00:41,081: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,083: Getting agent action\n", + "2023-10-10 21:00:41,085: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,086: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,088: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,090: Getting agent action\n", + "2023-10-10 21:00:41,092: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,093: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,095: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,097: Getting agent action\n", + "2023-10-10 21:00:41,099: Formatting agent action NODE_STARTUP\n", + "2023-10-10 21:00:41,100: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'startup']\n", + "2023-10-10 21:00:41,102: Initiating simulation step 24\n", + "2023-10-10 21:00:41,103: Stepping primaite session. Step counter: 25\n", + "2023-10-10 21:00:41,104: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,107: Getting agent action\n", + "2023-10-10 21:00:41,110: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,112: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,113: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,114: Getting agent action\n", + "2023-10-10 21:00:41,116: Formatting agent action NODE_OS_SCAN\n", + "2023-10-10 21:00:41,118: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", + "2023-10-10 21:00:41,120: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,122: Getting agent action\n", + "2023-10-10 21:00:41,124: Formatting agent action NODE_FOLDER_RESTORE\n", + "2023-10-10 21:00:41,126: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,127: Initiating simulation step 25\n", + "2023-10-10 21:00:41,129: Stepping primaite session. Step counter: 26\n", + "2023-10-10 21:00:41,130: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,132: Getting agent action\n", + "2023-10-10 21:00:41,134: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,136: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,137: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,139: Getting agent action\n", + "2023-10-10 21:00:41,141: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,142: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,144: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,145: Getting agent action\n", + "2023-10-10 21:00:41,147: Formatting agent action NODE_STARTUP\n", + "2023-10-10 21:00:41,148: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'startup']\n", + "2023-10-10 21:00:41,150: Initiating simulation step 26\n", + "2023-10-10 21:00:41,151: Stepping primaite session. Step counter: 27\n", + "2023-10-10 21:00:41,153: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,155: Getting agent action\n", + "2023-10-10 21:00:41,157: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,158: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,160: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,162: Getting agent action\n", + "2023-10-10 21:00:41,164: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,166: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,166: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,169: Getting agent action\n", + "2023-10-10 21:00:41,170: Formatting agent action NODE_SERVICE_START\n", + "2023-10-10 21:00:41,172: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,173: Initiating simulation step 27\n", + "2023-10-10 21:00:41,174: Stepping primaite session. Step counter: 28\n", + "2023-10-10 21:00:41,176: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,177: Getting agent action\n", + "2023-10-10 21:00:41,179: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,181: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,182: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,183: Getting agent action\n", + "2023-10-10 21:00:41,185: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,186: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,187: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,188: Getting agent action\n", + "2023-10-10 21:00:41,190: Formatting agent action NETWORK_NIC_DISABLE\n", + "2023-10-10 21:00:41,191: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,193: Initiating simulation step 28\n", + "2023-10-10 21:00:41,194: Stepping primaite session. Step counter: 29\n", + "2023-10-10 21:00:41,196: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,198: Getting agent action\n", + "2023-10-10 21:00:41,199: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,201: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,201: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,203: Getting agent action\n", + "2023-10-10 21:00:41,204: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,206: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,207: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,208: Getting agent action\n", + "2023-10-10 21:00:41,211: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:41,213: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 3]\n", + "2023-10-10 21:00:41,215: Initiating simulation step 29\n", + "2023-10-10 21:00:41,217: Stepping primaite session. Step counter: 30\n", + "2023-10-10 21:00:41,218: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,220: Getting agent action\n", + "2023-10-10 21:00:41,221: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,222: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,223: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,225: Getting agent action\n", + "2023-10-10 21:00:41,227: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,229: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,231: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,232: Getting agent action\n", + "2023-10-10 21:00:41,233: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:41,235: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:41,236: Initiating simulation step 30\n", + "2023-10-10 21:00:41,238: Stepping primaite session. Step counter: 31\n", + "2023-10-10 21:00:41,238: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,240: Getting agent action\n", + "2023-10-10 21:00:41,242: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,244: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,245: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,248: Getting agent action\n", + "2023-10-10 21:00:41,249: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,251: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,252: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,254: Getting agent action\n", + "2023-10-10 21:00:41,255: Formatting agent action NODE_FILE_RESTORE\n", + "2023-10-10 21:00:41,256: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,257: Initiating simulation step 31\n", + "2023-10-10 21:00:41,259: Stepping primaite session. Step counter: 32\n", + "2023-10-10 21:00:41,260: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,263: Getting agent action\n", + "2023-10-10 21:00:41,264: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,266: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,267: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,268: Getting agent action\n", + "2023-10-10 21:00:41,269: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,270: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,272: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,273: Getting agent action\n", + "2023-10-10 21:00:41,275: Formatting agent action NODE_FILE_RESTORE\n", + "2023-10-10 21:00:41,277: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,278: Initiating simulation step 32\n", + "2023-10-10 21:00:41,280: Stepping primaite session. Step counter: 33\n", + "2023-10-10 21:00:41,281: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,284: Getting agent action\n", + "2023-10-10 21:00:41,286: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,288: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,289: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,291: Getting agent action\n", + "2023-10-10 21:00:41,293: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,295: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,299: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,301: Getting agent action\n", + "2023-10-10 21:00:41,307: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:41,309: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 2]\n", + "2023-10-10 21:00:41,311: Initiating simulation step 33\n", + "2023-10-10 21:00:41,312: Stepping primaite session. Step counter: 34\n", + "2023-10-10 21:00:41,322: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,325: Getting agent action\n", + "2023-10-10 21:00:41,327: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,328: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,339: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,341: Getting agent action\n", + "2023-10-10 21:00:41,343: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,345: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,346: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,348: Getting agent action\n", + "2023-10-10 21:00:41,350: Formatting agent action NODE_FOLDER_CHECKHASH\n", + "2023-10-10 21:00:41,352: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,354: Initiating simulation step 34\n", + "2023-10-10 21:00:41,355: Stepping primaite session. Step counter: 35\n", + "2023-10-10 21:00:41,357: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,359: Getting agent action\n", + "2023-10-10 21:00:41,360: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,362: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,363: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,366: Getting agent action\n", + "2023-10-10 21:00:41,368: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,369: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,370: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,371: Getting agent action\n", + "2023-10-10 21:00:41,373: Formatting agent action NODE_FILE_RESTORE\n", + "2023-10-10 21:00:41,378: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,381: Initiating simulation step 35\n", + "2023-10-10 21:00:41,382: Stepping primaite session. Step counter: 36\n", + "2023-10-10 21:00:41,384: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,386: Getting agent action\n", + "2023-10-10 21:00:41,388: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,389: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,390: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,392: Getting agent action\n", + "2023-10-10 21:00:41,394: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,395: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,397: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,399: Getting agent action\n", + "2023-10-10 21:00:41,401: Formatting agent action NODE_SERVICE_RESUME\n", + "2023-10-10 21:00:41,402: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,403: Initiating simulation step 36\n", + "2023-10-10 21:00:41,405: Stepping primaite session. Step counter: 37\n", + "2023-10-10 21:00:41,406: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,408: Getting agent action\n", + "2023-10-10 21:00:41,409: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,410: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,413: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,415: Getting agent action\n", + "2023-10-10 21:00:41,416: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,417: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,418: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,420: Getting agent action\n", + "2023-10-10 21:00:41,422: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:41,422: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,424: Initiating simulation step 37\n", + "2023-10-10 21:00:41,425: Stepping primaite session. Step counter: 38\n", + "2023-10-10 21:00:41,427: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,429: Getting agent action\n", + "2023-10-10 21:00:41,431: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,432: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,433: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,434: Getting agent action\n", + "2023-10-10 21:00:41,436: Formatting agent action NODE_OS_SCAN\n", + "2023-10-10 21:00:41,438: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", + "2023-10-10 21:00:41,440: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,442: Getting agent action\n", + "2023-10-10 21:00:41,444: Formatting agent action NODE_FILE_CHECKHASH\n", + "2023-10-10 21:00:41,449: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,451: Initiating simulation step 38\n", + "2023-10-10 21:00:41,452: Stepping primaite session. Step counter: 39\n", + "2023-10-10 21:00:41,453: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,457: Getting agent action\n", + "2023-10-10 21:00:41,459: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,461: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,462: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,464: Getting agent action\n", + "2023-10-10 21:00:41,465: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,467: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,469: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,472: Getting agent action\n", + "2023-10-10 21:00:41,475: Formatting agent action NODE_SERVICE_SCAN\n", + "2023-10-10 21:00:41,476: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,477: Initiating simulation step 39\n", + "2023-10-10 21:00:41,479: Stepping primaite session. Step counter: 40\n", + "2023-10-10 21:00:41,480: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,482: Getting agent action\n", + "2023-10-10 21:00:41,484: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,486: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,487: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,489: Getting agent action\n", + "2023-10-10 21:00:41,490: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,491: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,493: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,495: Getting agent action\n", + "2023-10-10 21:00:41,497: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:41,498: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 7]\n", + "2023-10-10 21:00:41,499: Initiating simulation step 40\n", + "2023-10-10 21:00:41,501: Stepping primaite session. Step counter: 41\n", + "2023-10-10 21:00:41,503: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,504: Getting agent action\n", + "2023-10-10 21:00:41,506: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,507: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,509: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,510: Getting agent action\n", + "2023-10-10 21:00:41,513: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,514: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,515: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,517: Getting agent action\n", + "2023-10-10 21:00:41,519: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:41,520: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 0]\n", + "2023-10-10 21:00:41,521: Initiating simulation step 41\n", + "2023-10-10 21:00:41,522: Stepping primaite session. Step counter: 42\n", + "2023-10-10 21:00:41,523: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,524: Getting agent action\n", + "2023-10-10 21:00:41,527: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,528: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,530: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,532: Getting agent action\n", + "2023-10-10 21:00:41,534: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,535: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,536: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,538: Getting agent action\n", + "2023-10-10 21:00:41,540: Formatting agent action NODE_RESET\n", + "2023-10-10 21:00:41,541: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'reset']\n", + "2023-10-10 21:00:41,543: Initiating simulation step 42\n", + "2023-10-10 21:00:41,544: Stepping primaite session. Step counter: 43\n", + "2023-10-10 21:00:41,546: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,548: Getting agent action\n", + "2023-10-10 21:00:41,550: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,551: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,552: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,554: Getting agent action\n", + "2023-10-10 21:00:41,555: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,556: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,557: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,559: Getting agent action\n", + "2023-10-10 21:00:41,561: Formatting agent action NODE_FOLDER_RESTORE\n", + "2023-10-10 21:00:41,563: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,565: Initiating simulation step 43\n", + "2023-10-10 21:00:41,566: Stepping primaite session. Step counter: 44\n", + "2023-10-10 21:00:41,567: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,568: Getting agent action\n", + "2023-10-10 21:00:41,569: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,570: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,571: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,574: Getting agent action\n", + "2023-10-10 21:00:41,575: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,577: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,578: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,581: Getting agent action\n", + "2023-10-10 21:00:41,583: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:41,584: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:41,586: Initiating simulation step 44\n", + "2023-10-10 21:00:41,587: Stepping primaite session. Step counter: 45\n", + "2023-10-10 21:00:41,588: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,590: Getting agent action\n", + "2023-10-10 21:00:41,591: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,593: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,594: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,596: Getting agent action\n", + "2023-10-10 21:00:41,598: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,599: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,600: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,602: Getting agent action\n", + "2023-10-10 21:00:41,603: Formatting agent action NETWORK_ACL_REMOVERULE\n", + "2023-10-10 21:00:41,604: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 1]\n", + "2023-10-10 21:00:41,605: Initiating simulation step 45\n", + "2023-10-10 21:00:41,606: Stepping primaite session. Step counter: 46\n", + "2023-10-10 21:00:41,607: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,608: Getting agent action\n", + "2023-10-10 21:00:41,610: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,612: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,613: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,616: Getting agent action\n", + "2023-10-10 21:00:41,617: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,618: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,620: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,621: Getting agent action\n", + "2023-10-10 21:00:41,623: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:41,624: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", + "2023-10-10 21:00:41,626: Initiating simulation step 46\n", + "2023-10-10 21:00:41,627: Stepping primaite session. Step counter: 47\n", + "2023-10-10 21:00:41,628: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,630: Getting agent action\n", + "2023-10-10 21:00:41,632: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,634: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,635: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,636: Getting agent action\n", + "2023-10-10 21:00:41,638: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,639: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,640: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,642: Getting agent action\n", + "2023-10-10 21:00:41,643: Formatting agent action NETWORK_ACL_ADDRULE\n", + "2023-10-10 21:00:41,644: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,646: Initiating simulation step 47\n", + "2023-10-10 21:00:41,648: Stepping primaite session. Step counter: 48\n", + "2023-10-10 21:00:41,649: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,651: Getting agent action\n", + "2023-10-10 21:00:41,652: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,653: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,654: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,657: Getting agent action\n", + "2023-10-10 21:00:41,658: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,661: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,663: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,665: Getting agent action\n", + "2023-10-10 21:00:41,667: Formatting agent action NODE_SERVICE_STOP\n", + "2023-10-10 21:00:41,669: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,670: Initiating simulation step 48\n", + "2023-10-10 21:00:41,672: Stepping primaite session. Step counter: 49\n", + "2023-10-10 21:00:41,674: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,684: Getting agent action\n", + "2023-10-10 21:00:41,687: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,689: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,690: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,701: Getting agent action\n", + "2023-10-10 21:00:41,702: Formatting agent action NODE_FILE_CORRUPT\n", + "2023-10-10 21:00:41,705: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,715: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,718: Getting agent action\n", + "2023-10-10 21:00:41,720: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,734: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,737: Initiating simulation step 49\n", + "2023-10-10 21:00:41,738: Stepping primaite session. Step counter: 50\n", + "2023-10-10 21:00:41,740: Sending simulation state to agent client_1_green_user\n", + "2023-10-10 21:00:41,742: Getting agent action\n", + "2023-10-10 21:00:41,744: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,747: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,750: Sending simulation state to agent client_1_data_manipulation_red_bot\n", + "2023-10-10 21:00:41,753: Getting agent action\n", + "2023-10-10 21:00:41,755: Formatting agent action NODE_FILE_DELETE\n", + "2023-10-10 21:00:41,756: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,758: Sending simulation state to agent defender\n", + "2023-10-10 21:00:41,761: Getting agent action\n", + "2023-10-10 21:00:41,771: Formatting agent action DONOTHING\n", + "2023-10-10 21:00:41,774: Sending request to simulation: ['do_nothing']\n", + "2023-10-10 21:00:41,776: Initiating simulation step 50\n" + ] + } + ], "source": [ "for i in range(50):\n", " sess.step()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/sandbox.py b/sandbox.py new file mode 100644 index 00000000..8114c23a --- /dev/null +++ b/sandbox.py @@ -0,0 +1,19 @@ + +from primaite.game.session import PrimaiteSession + +from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS +import logging +_PRIMAITE_CONFIG['log_level']=logging.DEBUG +print(PRIMAITE_PATHS.app_log_dir_path) +import itertools +from primaite.game.session import PrimaiteSession +from primaite.simulator.sim_container import Simulation +from primaite.game.agent.interface import AbstractAgent +from primaite.simulator.network.networks import arcd_uc2_network +import yaml + +with open('example_config.yaml', 'r') as file: + cfg = yaml.safe_load(file) +sess = PrimaiteSession.from_config(cfg) + +sess.start_session() \ No newline at end of file diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index be80374b..73b9e998 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -6,7 +6,7 @@ from logging import Logger from typing import Dict, Final, List, Tuple, TYPE_CHECKING, Union import numpy as np -from gym import spaces +from gymnasium import spaces from primaite.acl.acl_rule import ACLRule from primaite.common.enums import FileSystemState, HardwareState, RulePermissionType, SoftwareState diff --git a/src/primaite/game/agent/GATE_agents.py b/src/primaite/game/agent/GATE_agents.py index ac1d776b..eb3c2987 100644 --- a/src/primaite/game/agent/GATE_agents.py +++ b/src/primaite/game/agent/GATE_agents.py @@ -1,57 +1,9 @@ -from primaite.game.agent.interface import AbstractGATEAgent -from arcd_gate.client.gate_client import GATEClient - - -class GATEMan(GATEClient): - - @property - def rl_framework(self) -> str: - return "SB3" - - @property - def rl_framework(self) -> str: - pass - - @property - def rl_algorithm(self) -> str: - pass - - @property - def seed(self) -> Optional[int]: - return None - - @property - def n_learn_episodes(self) -> int: - return 0 - - @property - def n_learn_steps(self) -> int: - return 0 - - @property - def n_eval_episodes(self) -> int: - return 0 - - @property - def n_eval_steps(self) -> int: - return 0 - - @property - def action_space(self) -> spaces.Space: - pass - - @property - def observation_space(self) -> spaces.Space: - pass - - def step(self, action: ActType) -> Tuple[np.ndarray, float, bool, bool, Dict]: - pass - - def reset(self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None) -> Tuple[np.ndarray, Dict]: - pass - - def close(self): - pass +from typing import Dict, Optional, Tuple +from primaite.game.agent.actions import ActionManager +from primaite.game.agent.interface import AbstractGATEAgent, ObsType +from primaite.game.agent.observations import ObservationSpace +from primaite.game.agent.rewards import RewardFunction +from gymnasium.core import ActType, ObsType class GATERLAgent(AbstractGATEAgent): ... @@ -61,4 +13,9 @@ class GATERLAgent(AbstractGATEAgent): # For example MultiAgentEnv in Ray allows sending a dict of observations of multiple agents, then it will reply # with the actions for those agents. + def __init__(self, agent_name: str | None, action_space: ActionManager | None, observation_space: ObservationSpace | None, reward_function: RewardFunction | None) -> None: + super().__init__(agent_name, action_space, observation_space, reward_function) + self.most_recent_action: ActType + def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: + return self.most_recent_action \ No newline at end of file diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index cba90305..6a1d5bcd 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -2,7 +2,7 @@ import itertools from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING -from gym import spaces +from gymnasium import spaces from primaite import getLogger from primaite.simulator.sim_container import Simulation diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 7b10f957..00f98f5c 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -1,10 +1,12 @@ from abc import ABC, abstractmethod from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING -from gym import spaces +from gymnasium import spaces from pydantic import BaseModel from primaite.simulator.sim_container import Simulation +from primaite import getLogger +_LOGGER = getLogger(__name__) if TYPE_CHECKING: from primaite.game.session import PrimaiteSession @@ -181,7 +183,7 @@ class LinkObservation(AbstractObservation): class FolderObservation(AbstractObservation): - def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = []) -> None: + def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder:int=2) -> None: """Initialise folder Observation, including files inside of the folder. :param where: Where in the simulation state dictionary to find the relevant information for this folder. @@ -203,6 +205,13 @@ class FolderObservation(AbstractObservation): self.where: Optional[Tuple[str]] = where self.files: List[FileObservation] = files + while len(self.files) < num_files_per_folder: + self.files.append(FileObservation()) + while len(self.files)> num_files_per_folder: + truncated_file = self.files.pop() + msg = f"Too many files in folde observation. Truncating file {truncated_file}" + _LOGGER.warn(msg) + raise UserWarning(msg) self.default_observation = { "health_status": 0, @@ -235,13 +244,13 @@ class FolderObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]], num_files_per_folder:int=2): where = parent_where + ["folders", config["folder_name"]] file_configs = config["files"] files = [FileObservation.from_config(config=f, session=session, parent_where=where) for f in file_configs] - return cls(where=where, files=files) + return cls(where=where, files=files, num_files_per_folder=num_files_per_folder) class NicObservation(AbstractObservation): @@ -277,7 +286,10 @@ class NodeObservation(AbstractObservation): folders: List[FolderObservation] = [], nics: List[NicObservation] = [], logon_status: bool = False, - ) -> None: + num_services_per_node: int = 2, + num_folders_per_node: int = 2, + num_files_per_folder: int = 2 + ) -> None: """ Configurable observation for a node in the simulation. @@ -302,7 +314,24 @@ class NodeObservation(AbstractObservation): self.where: Optional[Tuple[str]] = where self.services: List[ServiceObservation] = services + while len(self.services)num_services_per_node: + truncated_service = self.services.pop() + msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}" + _LOGGER.warn(msg) + raise UserWarning(msg) + # truncate service list + self.folders: List[FolderObservation] = folders + while len(self.folders) < num_folders_per_node: + # add an empty folder observation without `where` parameter that will always return default (blank) observations + self.folders.append(FolderObservation()) + while len(self.folders) > num_folders_per_node: + truncated_folder = self.folders.pop() + msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}" + self.nics: List[NicObservation] = nics self.logon_status: bool = logon_status @@ -349,7 +378,13 @@ class NodeObservation(AbstractObservation): @classmethod def from_config( - cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None + cls, + config: Dict, + session: "PrimaiteSession", + parent_where: Optional[List[str]] = None, + num_services_per_node: int = 2, + num_folders_per_node: int = 2, + num_files_per_folder: int = 2, ) -> "NodeObservation": node_uuid = session.ref_map_nodes[config["node_ref"]] if parent_where is None: @@ -360,12 +395,21 @@ class NodeObservation(AbstractObservation): svc_configs = config.get("services", {}) services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs] folder_configs = config.get("folders", {}) - folders = [FolderObservation.from_config(config=c, session=session, parent_where=where) for c in folder_configs] + folders = [FolderObservation.from_config(config=c, session=session, parent_where=where, num_files_per_folder=num_files_per_folder) for c in folder_configs] nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys() nic_configs = [{"nic_uuid": n for n in nic_uuids}] if nic_uuids else [] nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs] logon_status = config.get("logon_status", False) - return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status) + return cls( + where=where, + services=services, + folders=folders, + nics=nics, + logon_status=logon_status, + num_services_per_node = num_services_per_node, + num_folders_per_node = num_folders_per_node, + num_files_per_folder = num_files_per_folder, + ) class AclObservation(AbstractObservation): @@ -467,11 +511,12 @@ class AclObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation": node_ip_to_idx = {} - for node_idx, node_cfg in enumerate(config["node_order"]): - n_ref = node_cfg["node_ref"] - n_obj = session.simulation.network.nodes[session.ref_map_nodes[n_ref]] - for nic_uuid, nic_obj in n_obj.nics.items(): - node_ip_to_idx[nic_obj.ip_address] = node_idx + 2 + for ip_idx, ip_map_config in enumerate(config["ip_address_order"]): + node_ref = ip_map_config["node_ref"] + nic_num = ip_map_config["nic_num"] + node_obj = session.simulation.network.nodes[session.ref_map_nodes[node_ref]] + nic_obj = node_obj.ethernet_port[nic_num] + node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2 router_uuid = session.ref_map_nodes[config["router_node_ref"]] return cls( @@ -552,7 +597,16 @@ class UC2BlueObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession"): node_configs = config["nodes"] - nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs] + num_services_per_node = config["num_services_per_node"] + num_folders_per_node = config["num_folders_per_node"] + num_files_per_folder = config["num_files_per_folder"] + nodes = [NodeObservation.from_config( + config=n, + session=session, + num_services_per_node= num_services_per_node, + num_folders_per_node=num_folders_per_node, + num_files_per_folder=num_files_per_folder, + ) for n in node_configs] link_configs = config["links"] links = [LinkObservation.from_config(config=l, session=session) for l in link_configs] diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 7746f78c..d978b848 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -7,12 +7,16 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Optional, Tuple -from gymnasium.vector.utils import spaces +from gymnasium import spaces +from gymnasium.spaces.utils import flatten, flatten_space, unflatten +from gymnasium.core import ObsType, ActType + import numpy as np from pydantic import BaseModel from primaite import getLogger +from primaite.game.agent.GATE_agents import GATERLAgent from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent, RandomAgent from primaite.game.agent.observations import ( @@ -54,7 +58,7 @@ _LOGGER = getLogger(__name__) class PrimaiteGATEClient(GATEClient): def __init__(self, parent_session:"PrimaiteSession", service_port: int = 50000): super().__init__(service_port=service_port) - self.parent_session:"PrimaiteSession" + self.parent_session:"PrimaiteSession" = parent_session @property def rl_framework(self) -> str: @@ -86,21 +90,33 @@ class PrimaiteGATEClient(GATEClient): @property def action_space(self) -> spaces.Space: - return self.parent_session.rl_agent.action_space + return self.parent_session.rl_agent.action_space.space @property def observation_space(self) -> spaces.Space: - return self.parent_session.rl_agent.observation_space + print("YEEY0") + print(flatten_space(spaces.Dict({}))) + print("YEEY1") + # print(self.parent_session.rl_agent.observation_space.space) + return flatten_space(self.parent_session.rl_agent.observation_space.space) - def step(self, action: ActType) -> Tuple[ndarray, float, bool, bool, Dict]: + def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, Dict]: + self.parent_session.rl_agent.most_recent_action = action self.parent_session.step() - #TODO: not sure how to go about this. + obs = self.parent_session.rl_agent.observation_space.observe() + obs = flatten(self.parent_session.rl_agent.observation_space.space, obs) + rew = self.parent_session.rl_agent.reward_function.calculate() + term = False + trunc = False + info = {} + return obs, rew, term, trunc, info + def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ndarray, Dict]: - ... + self.parent_session.reset() def close(self): - ... + self.parent_session.close() class PrimaiteSessionOptions(BaseModel): ports: List[str] @@ -131,15 +147,11 @@ class PrimaiteSession: self.ref_map_nodes: Dict[str, Node] = {} self.ref_map_services: Dict[str, Service] = {} self.ref_map_links: Dict[str, Link] = {} + self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self) def start_session(self, opts="TODO..."): """Commence the session, this gives the gate client control over the simulation/agent loop.""" - ... - - def eval(self, opts="TODO..."): - ... - - + self.gate_client.start() def step(self): _LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}") @@ -172,12 +184,18 @@ class PrimaiteSession: # 10. primaite session receives the action from the agents and asks the simulation to apply each _LOGGER.debug(f"Sending request to simulation: {agent_request}") - self.simulation.apply_action(agent_request) + self.simulation.apply_request(agent_request) _LOGGER.debug(f"Initiating simulation step {self.step_counter}") self.simulation.apply_timestep(self.step_counter) self.step_counter += 1 + def reset(self): + pass + + def close(self): + pass + @classmethod def from_config(cls, cfg: dict) -> "PrimaiteSession": sess = cls() @@ -351,6 +369,7 @@ class PrimaiteSession: reward_function=rew_function, ) sess.agents.append(new_agent) + sess.rl_agent = new_agent elif agent_type == "RedDatabaseCorruptingAgent": new_agent = RandomAgent( agent_name=agent_cfg["ref"], diff --git a/src/primaite/transactions/transaction.py b/src/primaite/transactions/transaction.py index 7d5f747c..6b973ca3 100644 --- a/src/primaite/transactions/transaction.py +++ b/src/primaite/transactions/transaction.py @@ -7,7 +7,7 @@ from primaite.common.enums import AgentIdentifier if TYPE_CHECKING: import numpy as np - from gym import spaces + from gymnasium import spaces class Transaction(object): From b84ab84385a20efb50bee35b6145b0104d600fec Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 11 Oct 2023 14:08:55 +0100 Subject: [PATCH 21/35] Fix GATE Client to work successfully with GATE Server --- src/primaite/game/agent/observations.py | 6 +++--- src/primaite/game/session.py | 15 ++++++++------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 00f98f5c..a5a5fc77 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -487,7 +487,7 @@ class AclObservation(AbstractObservation): def space(self) -> spaces.Space: return spaces.Dict( { - "RULE": spaces.Dict( + "RULES": spaces.Dict( { i + 1: spaces.Dict( @@ -532,11 +532,11 @@ class NullObservation(AbstractObservation): self.default_observation: Dict = {} def observe(self, state: Dict) -> Dict: - return {} + return 0 @property def space(self) -> spaces.Space: - return spaces.Dict({}) + return spaces.Discrete(1) @classmethod def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation": diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index d978b848..ec6b8e86 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -94,26 +94,27 @@ class PrimaiteGATEClient(GATEClient): @property def observation_space(self) -> spaces.Space: - print("YEEY0") - print(flatten_space(spaces.Dict({}))) - print("YEEY1") - # print(self.parent_session.rl_agent.observation_space.space) return flatten_space(self.parent_session.rl_agent.observation_space.space) def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, Dict]: self.parent_session.rl_agent.most_recent_action = action self.parent_session.step() - obs = self.parent_session.rl_agent.observation_space.observe() + state = self.parent_session.simulation.describe_state() + obs = self.parent_session.rl_agent.observation_space.observe(state) obs = flatten(self.parent_session.rl_agent.observation_space.space, obs) - rew = self.parent_session.rl_agent.reward_function.calculate() + rew = self.parent_session.rl_agent.reward_function.calculate(state) term = False trunc = False info = {} return obs, rew, term, trunc, info - def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ndarray, Dict]: + def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ObsType, Dict]: self.parent_session.reset() + state = self.parent_session.simulation.describe_state() + obs = self.parent_session.rl_agent.observation_space.observe(state) + obs = flatten(self.parent_session.rl_agent.observation_space.space, obs) + return obs, {} def close(self): self.parent_session.close() From 70c1857bbc54414f48e42f065f791a8bec45caae Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 11 Oct 2023 15:49:41 +0100 Subject: [PATCH 22/35] Implement rewards for UC2 (draft) --- example_config.yaml | 13 ++- src/primaite/game/agent/observations.py | 30 +----- src/primaite/game/agent/rewards.py | 123 +++++++++++++++++++++--- src/primaite/game/agent/utils.py | 29 ++++++ 4 files changed, 152 insertions(+), 43 deletions(-) create mode 100644 src/primaite/game/agent/utils.py diff --git a/example_config.yaml b/example_config.yaml index a35c82e0..b700da5c 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -445,7 +445,18 @@ game_config: reward_function: reward_components: - - type: DUMMY + - type: DATABASE_FILE_INTEGRITY + weight: 0.5 + options: + node_ref: database_server + folder_name: database + file_name: database.db + - type: WEB_SERVER_404_PENALTY + weight: 0.5 + options: + node_ref: web_server + service_ref: web_server_database_client + agent_settings: # ... diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index a5a5fc77..ba1e8e66 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -5,41 +5,13 @@ from gymnasium import spaces from pydantic import BaseModel from primaite.simulator.sim_container import Simulation +from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite import getLogger _LOGGER = getLogger(__name__) if TYPE_CHECKING: from primaite.game.session import PrimaiteSession -NOT_PRESENT_IN_STATE = object() -""" -Need an object to return when the sim state does not contain a requested value. Cannot use None because sometimes -the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is a sentinel for this purpose. -""" - - -def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: - """ - Access an item from a deeply dictionary with a list of keys. - - For example, if the dictionary is {1: 'a', 2: {3: {4: 'b'}}}, then the key [2, 3, 4] would return 'b', and the key - [2, 3] would return {4: 'b'}. Raises a KeyError if specified key does not exist at any level of nesting. - - :param dictionary: Deeply nested dictionary - :type dictionary: Dict - :param keys: List of dict keys used to traverse the nested dict. Each item corresponds to one level of depth. - :type keys: List[Hashable] - :return: The value in the dictionary - :rtype: Any - """ - key_list = [*keys] # copy keys to a new list to prevent editing original list - if len(key_list) == 0: - return dictionary - k = key_list.pop(0) - if k not in dictionary: - return NOT_PRESENT_IN_STATE - return access_from_nested_dict(dictionary[k], key_list) - class AbstractObservation(ABC): @abstractmethod diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 18925edc..b7a4bb24 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -1,34 +1,131 @@ +from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE + from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple, TYPE_CHECKING +from primaite import getLogger +_LOGGER = getLogger(__name__) + +if TYPE_CHECKING: + from primaite.game.session import PrimaiteSession class AbstractReward: - def __init__(self): - ... @abstractmethod def calculate(self, state: Dict) -> float: - return 0.3 + return 0.0 + + @abstractmethod + @classmethod + def from_config(cls, config:dict) -> "AbstractReward": + return cls() class DummyReward(AbstractReward): def calculate(self, state: Dict) -> float: - return -0.1 + return 0.0 + + @classmethod + def from_config(cls, config: dict) -> "DummyReward": + return cls() + +class DatabaseFileIntegrity(AbstractReward): + def __init__(self, node_uuid:str, folder_name:str, file_name:str) -> None: + self.location_in_state = ["network", "node", node_uuid, "file_system", ""] + + def calculate(self, state: Dict) -> float: + database_file_state = access_from_nested_dict(state, self.location_in_state) + health_status = database_file_state['health_status'] + if health_status == "corrupted": + return -1 + elif health_status == "good": + return 1 + else: + return 0 + + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "DatabaseFileIntegrity": + node_ref = config.get("node_ref") + folder_name = config.get("folder_name") + file_name = config.get("file_name") + if not (node_ref): + _LOGGER.error(f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified") + return DummyReward() #TODO: better error handling + if not folder_name: + _LOGGER.error(f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified") + return DummyReward() # TODO: better error handling + if not file_name: + _LOGGER.error(f"{cls.__name__} could not be initialised from config because file_name parameter was not specified") + return DummyReward() # TODO: better error handling + node_uuid = session.ref_map_nodes[node_ref].uuid + if not node_uuid: + _LOGGER.error(f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation") + return DummyReward() # TODO: better error handling + + return cls(node_uuid = node_uuid, folder_name=folder_name, file_name=file_name) + +class WebServer404Penalty(AbstractReward): + def __init__(self, node_uuid:str, service_uuid:str) -> None: + self.location_in_state = ['network','node', node_uuid, 'services', service_uuid] + + def calculate(self, state: Dict) -> float: + web_service_state = access_from_nested_dict(state, self.location_in_state) + most_recent_return_code = web_service_state['most_recent_return_code'] + if most_recent_return_code == 200: + return 1 + elif most_recent_return_code == 404: + return -1 + else: + return 0 + + @classmethod + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "WebServer404Penalty": + node_ref = config.get("node_ref") + service_ref = config.get("service_ref") + if not (node_ref and service_ref): + msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config." + _LOGGER.warn(msg) + return DummyReward() #TODO: should we error out with incorrect inputs? Probably! + node_uuid = session.ref_map_nodes[node_ref].uuid + service_uuid = session.ref_map_services[service_ref].uuid + if not (node_uuid and service_uuid): + msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator." + _LOGGER.warn(msg) + return DummyReward() # TODO: consider erroring here as well + + return cls(node_uuid=node_uuid, service_uuid=service_uuid) class RewardFunction: - __rew_class_identifiers: Dict[str, type[AbstractReward]] = {"DUMMY": DummyReward} + __rew_class_identifiers: Dict[str, type[AbstractReward]] = { + "DUMMY": DummyReward, + "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity, + "WEB_SERVER_404_PENALTY": WebServer404Penalty, + } - def __init__(self, reward_function: AbstractReward): - self.reward: AbstractReward = reward_function + def __init__(self): + self.reward_components: List[Tuple[AbstractReward, float]] = [] + "attribute reward_components keeps track of reward components and the weights assigned to each." + + def regsiter_component(self, component:AbstractReward, weight:float=1.0) -> None: + self.reward_components.append((component, weight)) def calculate(self, state: Dict) -> float: - return self.reward.calculate(state) + total = 0.0 + for comp_and_weight in self.reward_components: + comp = comp_and_weight[0] + weight = comp_and_weight[1] + total += weight * comp.calculate(state=state) + return total @classmethod - def from_config(cls, cfg: Dict) -> "RewardFunction": - for rew_component_cfg in cfg["reward_components"]: + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "RewardFunction": + new = cls() + + for rew_component_cfg in config["reward_components"]: rew_type = rew_component_cfg["type"] - rew_component = cls.__rew_class_identifiers[rew_type]() - new = cls(reward_function=rew_component) + weight = rew_component_cfg["weight"] + rew_class = cls.__rew_class_identifiers[rew_type] + rew_instance = rew_class.from_config(config=rew_component_cfg.get('options',{}), session=session) + new.regsiter_component(component=rew_instance, weight=weight) return new diff --git a/src/primaite/game/agent/utils.py b/src/primaite/game/agent/utils.py new file mode 100644 index 00000000..ad6dbefe --- /dev/null +++ b/src/primaite/game/agent/utils.py @@ -0,0 +1,29 @@ +from typing import Dict, Sequence, Hashable, Any + +NOT_PRESENT_IN_STATE = object() +""" +Need an object to return when the sim state does not contain a requested value. Cannot use None because sometimes +the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is a sentinel for this purpose. +""" + +def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: + """ + Access an item from a deeply dictionary with a list of keys. + + For example, if the dictionary is {1: 'a', 2: {3: {4: 'b'}}}, then the key [2, 3, 4] would return 'b', and the key + [2, 3] would return {4: 'b'}. Raises a KeyError if specified key does not exist at any level of nesting. + + :param dictionary: Deeply nested dictionary + :type dictionary: Dict + :param keys: List of dict keys used to traverse the nested dict. Each item corresponds to one level of depth. + :type keys: List[Hashable] + :return: The value in the dictionary + :rtype: Any + """ + key_list = [*keys] # copy keys to a new list to prevent editing original list + if len(key_list) == 0: + return dictionary + k = key_list.pop(0) + if k not in dictionary: + return NOT_PRESENT_IN_STATE + return access_from_nested_dict(dictionary[k], key_list) \ No newline at end of file From 565af11dba80efbaa98897425f87678a0591b193 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 12 Oct 2023 09:59:45 +0100 Subject: [PATCH 23/35] Minor fixes to rewards --- src/primaite/game/agent/rewards.py | 16 ++++++++-------- src/primaite/game/session.py | 4 ++-- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index b7a4bb24..85da95da 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -15,9 +15,9 @@ class AbstractReward: def calculate(self, state: Dict) -> float: return 0.0 - @abstractmethod @classmethod - def from_config(cls, config:dict) -> "AbstractReward": + @abstractmethod + def from_config(cls, config:dict, session:"PrimaiteSession") -> "AbstractReward": return cls() @@ -26,12 +26,12 @@ class DummyReward(AbstractReward): return 0.0 @classmethod - def from_config(cls, config: dict) -> "DummyReward": + def from_config(cls, config: dict, session:"PrimaiteSession") -> "DummyReward": return cls() class DatabaseFileIntegrity(AbstractReward): def __init__(self, node_uuid:str, folder_name:str, file_name:str) -> None: - self.location_in_state = ["network", "node", node_uuid, "file_system", ""] + self.location_in_state = ["network", "nodes", node_uuid, "file_system", "folders",folder_name, "files", file_name] def calculate(self, state: Dict) -> float: database_file_state = access_from_nested_dict(state, self.location_in_state) @@ -57,7 +57,7 @@ class DatabaseFileIntegrity(AbstractReward): if not file_name: _LOGGER.error(f"{cls.__name__} could not be initialised from config because file_name parameter was not specified") return DummyReward() # TODO: better error handling - node_uuid = session.ref_map_nodes[node_ref].uuid + node_uuid = session.ref_map_nodes[node_ref] if not node_uuid: _LOGGER.error(f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation") return DummyReward() # TODO: better error handling @@ -66,7 +66,7 @@ class DatabaseFileIntegrity(AbstractReward): class WebServer404Penalty(AbstractReward): def __init__(self, node_uuid:str, service_uuid:str) -> None: - self.location_in_state = ['network','node', node_uuid, 'services', service_uuid] + self.location_in_state = ['network','nodes', node_uuid, 'services', service_uuid] def calculate(self, state: Dict) -> float: web_service_state = access_from_nested_dict(state, self.location_in_state) @@ -86,7 +86,7 @@ class WebServer404Penalty(AbstractReward): msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config." _LOGGER.warn(msg) return DummyReward() #TODO: should we error out with incorrect inputs? Probably! - node_uuid = session.ref_map_nodes[node_ref].uuid + node_uuid = session.ref_map_nodes[node_ref] service_uuid = session.ref_map_services[service_ref].uuid if not (node_uuid and service_uuid): msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator." @@ -124,7 +124,7 @@ class RewardFunction: for rew_component_cfg in config["reward_components"]: rew_type = rew_component_cfg["type"] - weight = rew_component_cfg["weight"] + weight = rew_component_cfg.get("weight",1.0) rew_class = cls.__rew_class_identifiers[rew_type] rew_instance = rew_class.from_config(config=rew_component_cfg.get('options',{}), session=session) new.regsiter_component(component=rew_instance, weight=weight) diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index ec6b8e86..406308b9 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -295,7 +295,7 @@ class PrimaiteSession: net.add_node(new_node) new_node.power_on() - sess.ref_map_nodes[node_ref] = new_node.uuid + sess.ref_map_nodes[node_ref] = new_node.uuid # TODO: fix incosistency with service and link. Node gets added by uuid, but service gets reference to object # 2. create links between nodes for link_cfg in links_cfg: @@ -350,7 +350,7 @@ class PrimaiteSession: action_space = ActionManager.from_config(sess, action_space_cfg) # CREATE REWARD FUNCTION - rew_function = RewardFunction.from_config(reward_function_cfg) + rew_function = RewardFunction.from_config(reward_function_cfg, session=sess) # CREATE AGENT if agent_type == "GreenWebBrowsingAgent": From e0f8c3c5eaf5bc8dea27d58af7160dfcf64d0b5f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 19 Oct 2023 01:56:40 +0100 Subject: [PATCH 24/35] Add documentation --- example_config.yaml | 10 +- src/primaite/game/__init__.py | 1 + src/primaite/game/agent/actions.py | 323 +++++++++++++++++++++++---- src/primaite/game/agent/interface.py | 4 +- src/primaite/game/agent/rewards.py | 74 +++--- src/primaite/game/session.py | 166 +++++++++----- 6 files changed, 445 insertions(+), 133 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index b700da5c..e16411fa 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -2,10 +2,10 @@ training_config: rl_framework: SB3 rl_algorithm: PPO seed: 333 - n_learn_episodes: 2 - n_learn_steps: 128 - n_eval_episodes: 2 - n_eval_steps: 128 + n_learn_episodes: 1 + n_learn_steps: 8 + n_eval_episodes: 0 + n_eval_steps: 8 game_config: @@ -451,6 +451,8 @@ game_config: node_ref: database_server folder_name: database file_name: database.db + + - type: WEB_SERVER_404_PENALTY weight: 0.5 options: diff --git a/src/primaite/game/__init__.py b/src/primaite/game/__init__.py index e69de29b..5d7a721f 100644 --- a/src/primaite/game/__init__.py +++ b/src/primaite/game/__init__.py @@ -0,0 +1 @@ +"""PrimAITE Game Layer.""" diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 6a1d5bcd..4c4aaab4 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1,6 +1,16 @@ +""" +This module contains the ActionManager class which belongs to the Agent class. + +An agent's action space is made up of a collection of actions. Each action is an instance of a subclass of +AbstractAction. The ActionManager is responsible for: + 1. Creating the action space from a list of action types. + 2. Converting an integer action choice into a specific action and parameter choice. + 3. Converting an action and parameter choice into a request which can be ingested by the PrimAITE simulation. This + ensures that requests conform to the simulator's request format. +""" import itertools from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING from gymnasium import spaces @@ -13,22 +23,9 @@ if TYPE_CHECKING: from primaite.game.session import PrimaiteSession -class ExecutionDefiniton(ABC): - """ - Converter from actions to simulator requests. - - Allows adding extra data/context that defines in more detail what an action means. - """ - - """ - Examples: - ('node', 'service', 'scan', 2, 0) means scan the first service on node index 2 - -> ['network', 'nodes', , 'services', , 'scan'w] - """ - ... - - class AbstractAction(ABC): + """Base class for actions.""" + @abstractmethod def __init__(self, manager: "ActionManager", **kwargs) -> None: """ @@ -46,6 +43,8 @@ class AbstractAction(ABC): """Dictionary describing the number of options for each parameter of this action. The keys of this dict must align with the keyword args of the form_request method.""" self.manager: ActionManager = manager + """Reference to the ActionManager which created this action. This is used to access the session and simulation + objects.""" @abstractmethod def form_request(self) -> List[str]: @@ -54,6 +53,8 @@ class AbstractAction(ABC): class DoNothingAction(AbstractAction): + """Action which does nothing. This is here to allow agents to be idle if they choose to.""" + def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) self.name = "DONOTHING" @@ -65,6 +66,7 @@ class DoNothingAction(AbstractAction): # with one option. This just aids the Action Manager to enumerate all possibilities. def form_request(self, **kwargs) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" return ["do_nothing"] @@ -77,12 +79,13 @@ class NodeServiceAbstractAction(AbstractAction): """ @abstractmethod - def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None: + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes, "service_id": num_services} self.verb: str def form_request(self, node_id: int, service_id: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_uuid = self.manager.get_node_uuid_by_idx(node_id) service_uuid = self.manager.get_service_uuid_by_idx(node_id, service_id) if node_uuid is None or service_uuid is None: @@ -91,54 +94,77 @@ class NodeServiceAbstractAction(AbstractAction): class NodeServiceScanAction(NodeServiceAbstractAction): + """Action which scans a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "scan" class NodeServiceStopAction(NodeServiceAbstractAction): + """Action which stops a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "stop" class NodeServiceStartAction(NodeServiceAbstractAction): + """Action which starts a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "start" class NodeServicePauseAction(NodeServiceAbstractAction): + """Action which pauses a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "pause" class NodeServiceResumeAction(NodeServiceAbstractAction): + """Action which resumes a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "resume" class NodeServiceRestartAction(NodeServiceAbstractAction): + """Action which restarts a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "restart" class NodeServiceDisableAction(NodeServiceAbstractAction): + """Action which disables a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "disable" class NodeServiceEnableAction(NodeServiceAbstractAction): + """Action which enables a service.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services) self.verb = "enable" class NodeFolderAbstractAction(AbstractAction): + """ + Base class for folder actions. + + Any action which applies to a folder and uses node_id and folder_id as its only two parameters can inherit from + this base class. + """ + @abstractmethod def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager=manager) @@ -146,6 +172,7 @@ class NodeFolderAbstractAction(AbstractAction): self.verb: str def form_request(self, node_id: int, folder_id: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_uuid = self.manager.get_node_uuid_by_idx(node_id) folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) if node_uuid is None or folder_uuid is None: @@ -154,30 +181,44 @@ class NodeFolderAbstractAction(AbstractAction): class NodeFolderScanAction(NodeFolderAbstractAction): + """Action which scans a folder.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "scan" class NodeFolderCheckhashAction(NodeFolderAbstractAction): + """Action which checks the hash of a folder.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "checkhash" class NodeFolderRepairAction(NodeFolderAbstractAction): + """Action which repairs a folder.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "repair" class NodeFolderRestoreAction(NodeFolderAbstractAction): + """Action which restores a folder.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "restore" class NodeFileAbstractAction(AbstractAction): + """Abstract base class for file actions. + + Any action which applies to a file and uses node_id, folder_id, and file_id as its only three parameters can inherit + from this base class. + """ + @abstractmethod def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager=manager) @@ -185,6 +226,7 @@ class NodeFileAbstractAction(AbstractAction): self.verb: str def form_request(self, node_id: int, folder_id: int, file_id: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_uuid = self.manager.get_node_uuid_by_idx(node_id) folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id) file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id) @@ -194,42 +236,60 @@ class NodeFileAbstractAction(AbstractAction): class NodeFileScanAction(NodeFileAbstractAction): + """Action which scans a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "scan" class NodeFileCheckhashAction(NodeFileAbstractAction): + """Action which checks the hash of a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "checkhash" class NodeFileDeleteAction(NodeFileAbstractAction): + """Action which deletes a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "delete" class NodeFileRepairAction(NodeFileAbstractAction): + """Action which repairs a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "repair" class NodeFileRestoreAction(NodeFileAbstractAction): + """Action which restores a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "restore" class NodeFileCorruptAction(NodeFileAbstractAction): + """Action which corrupts a file.""" + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs) self.verb = "corrupt" class NodeAbstractAction(AbstractAction): + """ + Abstract base class for node actions. + + Any action which applies to a node and uses node_id as its only parameter can inherit from this base class. + """ + @abstractmethod def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager) @@ -237,35 +297,46 @@ class NodeAbstractAction(AbstractAction): self.verb: str def form_request(self, node_id: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_uuid = self.manager.get_node_uuid_by_idx(node_id) return ["network", "node", node_uuid, self.verb] class NodeOSScanAction(NodeAbstractAction): + """Action which scans a node's OS.""" + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "scan" class NodeShutdownAction(NodeAbstractAction): + """Action which shuts down a node.""" + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "shutdown" class NodeStartupAction(NodeAbstractAction): + """Action which starts up a node.""" + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "startup" class NodeResetAction(NodeAbstractAction): + """Action which resets a node.""" + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes) self.verb = "reset" class NetworkACLAddRuleAction(AbstractAction): + """Action which adds a rule to a router's ACL.""" + def __init__( self, manager: "ActionManager", @@ -276,6 +347,21 @@ class NetworkACLAddRuleAction(AbstractAction): num_protocols: int, **kwargs, ) -> None: + """Init method for NetworkACLAddRuleAction. + + :param manager: Reference to the ActionManager which created this action. + :type manager: ActionManager + :param target_router_uuid: UUID of the router to which the ACL rule should be added. + :type target_router_uuid: str + :param max_acl_rules: Maximum number of ACL rules that can be added to the router. + :type max_acl_rules: int + :param num_ips: Number of IP addresses in the simulation. + :type num_ips: int + :param num_ports: Number of ports in the simulation. + :type num_ports: int + :param num_protocols: Number of protocols in the simulation. + :type num_protocols: int + """ super().__init__(manager=manager) num_permissions = 3 self.shape: Dict[str, int] = { @@ -290,8 +376,16 @@ class NetworkACLAddRuleAction(AbstractAction): self.target_router_uuid: str = target_router_uuid def form_request( - self, position, permission, source_ip_id, dest_ip_id, source_port_id, dest_port_id, protocol_id + self, + position: int, + permission: int, + source_ip_id: int, + dest_ip_id: int, + source_port_id: int, + dest_port_id: int, + protocol_id: int, ) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" if permission == 0: permission_str = "UNUSED" return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS @@ -354,22 +448,51 @@ class NetworkACLAddRuleAction(AbstractAction): class NetworkACLRemoveRuleAction(AbstractAction): + """Action which removes a rule from a router's ACL.""" + def __init__(self, manager: "ActionManager", target_router_uuid: str, max_acl_rules: int, **kwargs) -> None: + """Init method for NetworkACLRemoveRuleAction. + + :param manager: Reference to the ActionManager which created this action. + :type manager: ActionManager + :param target_router_uuid: UUID of the router from which the ACL rule should be removed. + :type target_router_uuid: str + :param max_acl_rules: Maximum number of ACL rules that can be added to the router. + :type max_acl_rules: int + """ super().__init__(manager=manager) self.shape: Dict[str, int] = {"position": max_acl_rules} self.target_router_uuid: str = target_router_uuid def form_request(self, position: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position] class NetworkNICAbstractAction(AbstractAction): + """ + Abstract base class for NIC actions. + + Any action which applies to a NIC and uses node_id and nic_id as its only two parameters can inherit from this base + class. + """ + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: + """Init method for NetworkNICAbstractAction. + + :param manager: Reference to the ActionManager which created this action. + :type manager: ActionManager + :param num_nodes: Number of nodes in the simulation. + :type num_nodes: int + :param max_nics_per_node: Maximum number of NICs per node. + :type max_nics_per_node: int + """ super().__init__(manager=manager) self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} self.verb: str def form_request(self, node_id: int, nic_id: int) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" node_uuid = self.manager.get_node_uuid_by_idx(node_idx=node_id) nic_uuid = self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id) if node_uuid is None or nic_uuid is None: @@ -385,45 +508,24 @@ class NetworkNICAbstractAction(AbstractAction): class NetworkNICEnableAction(NetworkNICAbstractAction): + """Action which enables a NIC.""" + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) self.verb = "enable" class NetworkNICDisableAction(NetworkNICAbstractAction): + """Action which disables a NIC.""" + def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs) self.verb = "disable" -# class NetworkNICDisableAction(AbstractAction): -# def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None: -# super().__init__(manager=manager) -# self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node} - -# def form_request(self, node_id: int, nic_id: int) -> List[str]: -# return [ -# "network", -# "node", -# self.manager.get_node_uuid_by_idx(node_idx=node_id), -# "nic", -# self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id), -# "disable", -# ] - - class ActionManager: - # let the action manager handle the conversion of action spaces into a single discrete integer space. - # + """Class which manages the action space for an agent.""" - # when action space is created, it will take subspaces and generate an action map by enumerating all possibilities, - # BUT, the action map can be provided in the config, in which case it will use that. - - # action map is basically just a mapping between integer and CAOS action (incl. parameter values) - # for example the action map can be: - # 0: DONOTHING - # 1: NODE, FILE, SCAN, NODEID=2, FOLDERID=1, FILEID=0 - # 2: ...... __act_class_identifiers: Dict[str, type] = { "DONOTHING": DoNothingAction, "NODE_SERVICE_SCAN": NodeServiceScanAction, @@ -453,6 +555,7 @@ class ActionManager: "NETWORK_NIC_ENABLE": NetworkNICEnableAction, "NETWORK_NIC_DISABLE": NetworkNICDisableAction, } + """Dictionary which maps action type strings to the corresponding action class.""" def __init__( self, @@ -469,6 +572,33 @@ class ActionManager: ip_address_list: Optional[List[str]] = None, # to allow us to map an index to an ip address. act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions ) -> None: + """Init method for ActionManager. + + :param session: Reference to the session to which the agent belongs. + :type session: PrimaiteSession + :param actions: List of action types which should be made available to the agent. + :type actions: List[str] + :param node_uuids: List of node UUIDs that this agent can act on. + :type node_uuids: List[str] + :param max_folders_per_node: Maximum number of folders per node. Used for calculating action shape. + :type max_folders_per_node: int + :param max_files_per_folder: Maximum number of files per folder. Used for calculating action shape. + :type max_files_per_folder: int + :param max_services_per_node: Maximum number of services per node. Used for calculating action shape. + :type max_services_per_node: int + :param max_nics_per_node: Maximum number of NICs per node. Used for calculating action shape. + :type max_nics_per_node: int + :param max_acl_rules: Maximum number of ACL rules per router. Used for calculating action shape. + :type max_acl_rules: int + :param protocols: List of protocols that are available in the simulation. Used for calculating action shape. + :type protocols: List[str] + :param ports: List of ports that are available in the simulation. Used for calculating action shape. + :type ports: List[str] + :param ip_address_list: List of IP addresses that known to this agent. Used for calculating action shape. + :type ip_address_list: Optional[List[str]] + :param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions. + :type act_map: Optional[Dict[int, Dict]] + """ self.session: "PrimaiteSession" = session self.sim: Simulation = self.session.simulation self.node_uuids: List[str] = node_uuids @@ -578,18 +708,48 @@ class ActionManager: @property def space(self) -> spaces.Space: + """Return the gymnasium action space for this agent.""" return spaces.Discrete(len(self.action_map)) - def get_node_uuid_by_idx(self, node_idx): + def get_node_uuid_by_idx(self, node_idx: int) -> str: + """Get the node UUID corresponding to the given index. + + :param node_idx: The index of the node to retrieve. + :type node_idx: int + :return: The node UUID. + :rtype: str + """ return self.node_uuids[node_idx] - def get_folder_uuid_by_idx(self, node_idx, folder_idx) -> Optional[str]: + def get_folder_uuid_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]: + """Get the folder UUID corresponding to the given node and folder indices. + + :param node_idx: The index of the node. + :type node_idx: int + :param folder_idx: The index of the folder on the node. + :type folder_idx: int + :return: The UUID of the folder. Or None if the node has fewer folders than the given index. + :rtype: Optional[str] + """ + node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] folder_uuids = list(node.file_system.folders.keys()) return folder_uuids[folder_idx] if len(folder_uuids) > folder_idx else None - def get_file_uuid_by_idx(self, node_idx, folder_idx, file_idx) -> Optional[str]: + def get_file_uuid_by_idx(self, node_idx: int, folder_idx: int, file_idx: int) -> Optional[str]: + """Get the file UUID corresponding to the given node, folder, and file indices. + + :param node_idx: The index of the node. + :type node_idx: int + :param folder_idx: The index of the folder on the node. + :type folder_idx: int + :param file_idx: The index of the file in the folder. + :type file_idx: int + :return: The UUID of the file. Or None if the node has fewer folders than the given index, or the folder has + fewer files than the given index. + :rtype: Optional[str] + """ node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] folder_uuids = list(node.file_system.folders.keys()) @@ -599,22 +759,64 @@ class ActionManager: file_uuids = list(folder.files.keys()) return file_uuids[file_idx] if len(file_uuids) > file_idx else None - def get_service_uuid_by_idx(self, node_idx, service_idx) -> Optional[str]: + def get_service_uuid_by_idx(self, node_idx: int, service_idx: int) -> Optional[str]: + """Get the service UUID corresponding to the given node and service indices. + + :param node_idx: The index of the node. + :type node_idx: int + :param service_idx: The index of the service on the node. + :type service_idx: int + :return: The UUID of the service. Or None if the node has fewer services than the given index. + :rtype: Optional[str] + """ node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] service_uuids = list(node.services.keys()) return service_uuids[service_idx] if len(service_uuids) > service_idx else None def get_internet_protocol_by_idx(self, protocol_idx: int) -> str: + """Get the internet protocol corresponding to the given index. + + :param protocol_idx: The index of the protocol to retrieve. + :type protocol_idx: int + :return: The protocol. + :rtype: str + """ return self.protocols[protocol_idx] def get_ip_address_by_idx(self, ip_idx: int) -> str: + """ + Get the IP address corresponding to the given index. + + :param ip_idx: The index of the IP address to retrieve. + :type ip_idx: int + :return: The IP address. + :rtype: str + """ return self.ip_address_list[ip_idx] def get_port_by_idx(self, port_idx: int) -> str: + """ + Get the port corresponding to the given index. + + :param port_idx: The index of the port to retrieve. + :type port_idx: int + :return: The port. + :rtype: str + """ return self.ports[port_idx] def get_nic_uuid_by_idx(self, node_idx: int, nic_idx: int) -> str: + """ + Get the NIC UUID corresponding to the given node and NIC indices. + + :param node_idx: The index of the node. + :type node_idx: int + :param nic_idx: The index of the NIC on the node. + :type nic_idx: int + :return: The NIC UUID. + :rtype: str + """ node_uuid = self.get_node_uuid_by_idx(node_idx) node_obj = self.sim.network.nodes[node_uuid] nics = list(node_obj.nics.keys()) @@ -624,6 +826,31 @@ class ActionManager: @classmethod def from_config(cls, session: "PrimaiteSession", cfg: Dict) -> "ActionManager": + """ + Construct an ActionManager from a config definition. + + The action space config supports the following three sections: + 1. ``action_list`` + ``action_list`` contians a list action components which need to be included in the action space. + Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options + which will be passed to the action class's __init__ method during initialisation. + 2. ``action_map`` + Since the agent uses a discrete action space which acts as a flattened version of the component-based + action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful + action and values of parameters. For example action 0 can correspond to do nothing, action 1 can + correspond to "NODE_SERVICE_SCAN" with ``node_id=1`` and ``service_id=1``, action 2 can be " + 3. ``options`` + ``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method. + These options are used to calculate the shape of the action space, and to provide additional information + to the ActionManager which is required to convert the agent's action choice into a CAOS request. + + :param session: The Primaite Session to which the agent belongs. + :type session: PrimaiteSession + :param cfg: The action space config. + :type cfg: Dict + :return: The constructed ActionManager. + :rtype: ActionManager + """ obj = cls( session=session, actions=cfg["action_list"], diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 817e59b1..5f121fcc 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -1,6 +1,4 @@ -# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' -# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word -# If you disagree, make a comment in the PR review and we can discuss +"""Interface for agents.""" from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple, TypeAlias, Union diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 85da95da..67e6ee50 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -1,8 +1,9 @@ -from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE - from abc import ABC, abstractmethod from typing import Any, Dict, List, Tuple, TYPE_CHECKING + from primaite import getLogger +from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE + _LOGGER = getLogger(__name__) if TYPE_CHECKING: @@ -10,14 +11,13 @@ if TYPE_CHECKING: class AbstractReward: - @abstractmethod def calculate(self, state: Dict) -> float: return 0.0 @classmethod @abstractmethod - def from_config(cls, config:dict, session:"PrimaiteSession") -> "AbstractReward": + def from_config(cls, config: dict, session: "PrimaiteSession") -> "AbstractReward": return cls() @@ -26,16 +26,26 @@ class DummyReward(AbstractReward): return 0.0 @classmethod - def from_config(cls, config: dict, session:"PrimaiteSession") -> "DummyReward": + def from_config(cls, config: dict, session: "PrimaiteSession") -> "DummyReward": return cls() + class DatabaseFileIntegrity(AbstractReward): - def __init__(self, node_uuid:str, folder_name:str, file_name:str) -> None: - self.location_in_state = ["network", "nodes", node_uuid, "file_system", "folders",folder_name, "files", file_name] + def __init__(self, node_uuid: str, folder_name: str, file_name: str) -> None: + self.location_in_state = [ + "network", + "nodes", + node_uuid, + "file_system", + "folders", + folder_name, + "files", + file_name, + ] def calculate(self, state: Dict) -> float: database_file_state = access_from_nested_dict(state, self.location_in_state) - health_status = database_file_state['health_status'] + health_status = database_file_state["health_status"] if health_status == "corrupted": return -1 elif health_status == "good": @@ -48,29 +58,39 @@ class DatabaseFileIntegrity(AbstractReward): node_ref = config.get("node_ref") folder_name = config.get("folder_name") file_name = config.get("file_name") - if not (node_ref): - _LOGGER.error(f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified") - return DummyReward() #TODO: better error handling + if not node_ref: + _LOGGER.error( + f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified" + ) + return DummyReward() # TODO: better error handling if not folder_name: - _LOGGER.error(f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified") - return DummyReward() # TODO: better error handling + _LOGGER.error( + f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified" + ) + return DummyReward() # TODO: better error handling if not file_name: - _LOGGER.error(f"{cls.__name__} could not be initialised from config because file_name parameter was not specified") - return DummyReward() # TODO: better error handling + _LOGGER.error( + f"{cls.__name__} could not be initialised from config because file_name parameter was not specified" + ) + return DummyReward() # TODO: better error handling node_uuid = session.ref_map_nodes[node_ref] if not node_uuid: - _LOGGER.error(f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation") - return DummyReward() # TODO: better error handling + _LOGGER.error( + f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation" + ) + return DummyReward() # TODO: better error handling + + return cls(node_uuid=node_uuid, folder_name=folder_name, file_name=file_name) - return cls(node_uuid = node_uuid, folder_name=folder_name, file_name=file_name) class WebServer404Penalty(AbstractReward): - def __init__(self, node_uuid:str, service_uuid:str) -> None: - self.location_in_state = ['network','nodes', node_uuid, 'services', service_uuid] + def __init__(self, node_uuid: str, service_uuid: str) -> None: + self.location_in_state = ["network", "nodes", node_uuid, "services", service_uuid] def calculate(self, state: Dict) -> float: web_service_state = access_from_nested_dict(state, self.location_in_state) - most_recent_return_code = web_service_state['most_recent_return_code'] + most_recent_return_code = web_service_state["most_recent_return_code"] + # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan. if most_recent_return_code == 200: return 1 elif most_recent_return_code == 404: @@ -85,13 +105,13 @@ class WebServer404Penalty(AbstractReward): if not (node_ref and service_ref): msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config." _LOGGER.warn(msg) - return DummyReward() #TODO: should we error out with incorrect inputs? Probably! + return DummyReward() # TODO: should we error out with incorrect inputs? Probably! node_uuid = session.ref_map_nodes[node_ref] service_uuid = session.ref_map_services[service_ref].uuid if not (node_uuid and service_uuid): msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator." _LOGGER.warn(msg) - return DummyReward() # TODO: consider erroring here as well + return DummyReward() # TODO: consider erroring here as well return cls(node_uuid=node_uuid, service_uuid=service_uuid) @@ -101,13 +121,13 @@ class RewardFunction: "DUMMY": DummyReward, "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity, "WEB_SERVER_404_PENALTY": WebServer404Penalty, - } + } def __init__(self): self.reward_components: List[Tuple[AbstractReward, float]] = [] "attribute reward_components keeps track of reward components and the weights assigned to each." - def regsiter_component(self, component:AbstractReward, weight:float=1.0) -> None: + def regsiter_component(self, component: AbstractReward, weight: float = 1.0) -> None: self.reward_components.append((component, weight)) def calculate(self, state: Dict) -> float: @@ -124,8 +144,8 @@ class RewardFunction: for rew_component_cfg in config["reward_components"]: rew_type = rew_component_cfg["type"] - weight = rew_component_cfg.get("weight",1.0) + weight = rew_component_cfg.get("weight", 1.0) rew_class = cls.__rew_class_identifiers[rew_type] - rew_instance = rew_class.from_config(config=rew_component_cfg.get('options',{}), session=session) + rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}), session=session) new.regsiter_component(component=rew_instance, weight=weight) return new diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 406308b9..f29d03dd 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -1,39 +1,17 @@ -# What do? Be an entry point for using PrimAITE -# 1. parse monoconfig -# 2. craete simulation -# 3. create actors and configure their actions/observations/rewards/ anything else -# 4. Create connection with ARCD GATE -# 5. idk - +"""PrimAITE session - the main entry point to training agents on PrimAITE.""" from ipaddress import IPv4Address from typing import Any, Dict, List, Optional, Tuple + +from arcd_gate.client.gate_client import ActType, GATEClient from gymnasium import spaces -from gymnasium.spaces.utils import flatten, flatten_space, unflatten -from gymnasium.core import ObsType, ActType - -import numpy as np - +from gymnasium.core import ActType, ObsType +from gymnasium.spaces.utils import flatten, flatten_space from pydantic import BaseModel from primaite import getLogger -from primaite.game.agent.GATE_agents import GATERLAgent from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent, RandomAgent -from primaite.game.agent.observations import ( - AclObservation, - FileObservation, - FolderObservation, - ICSObservation, - LinkObservation, - NicObservation, - NodeObservation, - NullObservation, - ObservationSpace, - ServiceObservation, - UC2BlueObservation, - UC2GreenObservation, - UC2RedObservation, -) +from primaite.game.agent.observations import ObservationSpace from primaite.game.agent.rewards import RewardFunction from primaite.simulator.network.hardware.base import Link, NIC, Node from primaite.simulator.network.hardware.nodes.computer import Computer @@ -50,53 +28,75 @@ from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.service import Service -from arcd_gate.client.gate_client import GATEClient, ActType -from numpy import ndarray - _LOGGER = getLogger(__name__) + class PrimaiteGATEClient(GATEClient): - def __init__(self, parent_session:"PrimaiteSession", service_port: int = 50000): + def __init__(self, parent_session: "PrimaiteSession", service_port: int = 50000): + """Create a new GATE client for PrimAITE. + + :param parent_session: The parent session object. + :type parent_session: PrimaiteSession + :param service_port: The port on which the GATE service is running. + :type service_port: int, optional""" super().__init__(service_port=service_port) - self.parent_session:"PrimaiteSession" = parent_session + self.parent_session: "PrimaiteSession" = parent_session @property def rl_framework(self) -> str: + """The reinforcement learning framework to use.""" return self.parent_session.training_options.rl_framework @property def rl_algorithm(self) -> str: + """The reinforcement learning algorithm to use.""" return self.parent_session.training_options.rl_algorithm @property def seed(self) -> int | None: + """The seed to use for the environment's random number generator.""" return self.parent_session.training_options.seed @property def n_learn_episodes(self) -> int: + """The number of episodes in each learning run.""" return self.parent_session.training_options.n_learn_episodes @property def n_learn_steps(self) -> int: + """The number of steps in each learning episode.""" return self.parent_session.training_options.n_learn_steps @property def n_eval_episodes(self) -> int: + """The number of episodes in each evaluation run.""" return self.parent_session.training_options.n_eval_episodes @property def n_eval_steps(self) -> int: + """The number of steps in each evaluation episode.""" return self.parent_session.training_options.n_eval_steps @property def action_space(self) -> spaces.Space: + """The gym action space of the agent.""" return self.parent_session.rl_agent.action_space.space @property def observation_space(self) -> spaces.Space: + """The gymnasium observation space of the agent.""" return flatten_space(self.parent_session.rl_agent.observation_space.space) def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, Dict]: + """Take a step in the environment. + + This method is called by GATE to advance the simulation by one timestep. + + :param action: The agent's action. + :type action: ActType + :return: The observation, reward, terminal flag, truncated flag, and info dictionary. + :rtype: Tuple[ObsType, float, bool, bool, Dict] + """ self.parent_session.rl_agent.most_recent_action = action self.parent_session.step() state = self.parent_session.simulation.describe_state() @@ -108,8 +108,19 @@ class PrimaiteGATEClient(GATEClient): info = {} return obs, rew, term, trunc, info - def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ObsType, Dict]: + """Reset the environment. + + This method is called when the environment is initialized and at the end of each episode. + + :param seed: The seed to use for the environment's random number generator. + :type seed: int, optional + :param options: Additional options for the reset. None are used by PrimAITE but this is included for + compatibility with GATE. + :type options: dict[str, Any], optional + :return: The initial observation and an empty info dictionary. + :rtype: Tuple[ObsType, Dict] + """ self.parent_session.reset() state = self.parent_session.simulation.describe_state() obs = self.parent_session.rl_agent.observation_space.observe(state) @@ -117,44 +128,78 @@ class PrimaiteGATEClient(GATEClient): return obs, {} def close(self): + """Close the session, this will stop the gate client and close the simulation.""" self.parent_session.close() + class PrimaiteSessionOptions(BaseModel): + """Global options which are applicable to all of the agents in the game. + + Currently this is used to restrict which ports and protocols exist in the world of the simulation.""" + ports: List[str] protocols: List[str] -class TrainingOptions(BaseModel): - rl_framework:str - rl_algorithm:str - seed:Optional[int] - n_learn_episodes:int - n_learn_steps:int - n_eval_episodes:int - n_eval_steps:int +class TrainingOptions(BaseModel): + """Options for training the RL agent.""" + + rl_framework: str + rl_algorithm: str + seed: Optional[int] + n_learn_episodes: int + n_learn_steps: int + n_eval_episodes: int + n_eval_steps: int class PrimaiteSession: + """ + The main entrypoint for PrimAITE sessions, this coordinates a simulation, agents, and connections to ARCD GATE. + """ + def __init__(self): self.simulation: Simulation = Simulation() + """Simulation object with which the agents will interact.""" self.agents: List[AbstractAgent] = [] + """List of agents.""" self.rl_agent: AbstractAgent - # which of the agents should be used for sending RL data to GATE client? + """The agent from the list which communicates with GATE to perform reinforcement learning.""" self.step_counter: int = 0 + """Current timestep within the episode.""" self.episode_counter: int = 0 + """Current episode number.""" self.options: PrimaiteSessionOptions + """Special options that apply for the entire game.""" self.training_options: TrainingOptions + """Options specific to agent training.""" self.ref_map_nodes: Dict[str, Node] = {} + """Mapping from unique node reference name to node object. Used when parsing config files.""" self.ref_map_services: Dict[str, Service] = {} + """Mapping from human-readable service reference to service object. Used for parsing config files.""" self.ref_map_links: Dict[str, Link] = {} + """Mapping from human-readable link reference to link object. Used when parsing config files.""" self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self) + """Reference to a GATE Client object, which will send data to GATE service for training RL agent.""" def start_session(self, opts="TODO..."): - """Commence the session, this gives the gate client control over the simulation/agent loop.""" + """Commence the training session, this gives the GATE client control over the simulation/agent loop.""" self.gate_client.start() def step(self): + """ + Perform one step of the simulation/agent loop. + + This is the main loop of the game. It corresponds to one timestep in the simulation, and one action from each + agent. The steps are as follows: + 1. The simulation state is updated. + 2. The simulation state is sent to each agent. + 3. Each agent converts the state to an observation and calculates a reward. + 4. Each agent chooses an action based on the observation. + 5. Each agent converts the action to a request. + 6. The simulation applies the requests. + """ _LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}") # currently designed with assumption that all agents act once per step in order @@ -192,19 +237,36 @@ class PrimaiteSession: self.step_counter += 1 def reset(self): - pass + """Reset the session, this will reset the simulation.""" + return NotImplemented def close(self): - pass + """Close the session, this will stop the gate client and close the simulation.""" + return NotImplemented @classmethod def from_config(cls, cfg: dict) -> "PrimaiteSession": + """Create a PrimaiteSession object from a config dictionary. + + The config dictionary should have the following top-level keys: + 1. training_config: options for training the RL agent. Used by GATE. + 2. game_config: options for the game itself. Used by PrimaiteSession. + 3. simulation: defines the network topology and the initial state of the simulation. + + The specification for each of the three major areas is described in a separate documentation page. + # TODO: create documentation page and add links to it here. + + :param cfg: The config dictionary. + :type cfg: dict + :return: A PrimaiteSession object. + :rtype: PrimaiteSession + """ sess = cls() sess.options = PrimaiteSessionOptions( ports=cfg["game_config"]["ports"], protocols=cfg["game_config"]["protocols"], ) - sess.training_options = TrainingOptions(**cfg['training_config']) + sess.training_options = TrainingOptions(**cfg["training_config"]) sim = sess.simulation net = sim.network @@ -295,7 +357,11 @@ class PrimaiteSession: net.add_node(new_node) new_node.power_on() - sess.ref_map_nodes[node_ref] = new_node.uuid # TODO: fix incosistency with service and link. Node gets added by uuid, but service gets reference to object + sess.ref_map_nodes[ + node_ref + ] = ( + new_node.uuid + ) # TODO: fix incosistency with service and link. Node gets added by uuid, but service by object # 2. create links between nodes for link_cfg in links_cfg: @@ -314,12 +380,10 @@ class PrimaiteSession: # 3. create agents game_cfg = cfg["game_config"] - ports_cfg = game_cfg["ports"] - protocols_cfg = game_cfg["protocols"] agents_cfg = game_cfg["agents"] for agent_cfg in agents_cfg: - agent_ref = agent_cfg["ref"] + agent_ref = agent_cfg["ref"] # noqa: F841 agent_type = agent_cfg["type"] action_space_cfg = agent_cfg["action_space"] observation_space_cfg = agent_cfg["observation_space"] From 49e78d529114011ce6926f516948ad6978ae194b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 19 Oct 2023 09:36:23 +0100 Subject: [PATCH 25/35] Add docstrings and fix formatting on many things --- sandbox.py | 23 +- src/primaite/game/agent/GATE_agents.py | 16 +- src/primaite/game/agent/actions.py | 13 +- src/primaite/game/agent/interface.py | 45 ++- src/primaite/game/agent/observations.py | 410 ++++++++++++++++++--- src/primaite/game/agent/rewards.py | 142 ++++++- src/primaite/game/agent/scripted_agents.py | 9 +- src/primaite/game/agent/utils.py | 5 +- src/primaite/game/session.py | 28 +- 9 files changed, 600 insertions(+), 91 deletions(-) diff --git a/sandbox.py b/sandbox.py index 8114c23a..ab5e701f 100644 --- a/sandbox.py +++ b/sandbox.py @@ -1,19 +1,22 @@ - -from primaite.game.session import PrimaiteSession +# flake8: noqa +import logging from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS -import logging -_PRIMAITE_CONFIG['log_level']=logging.DEBUG +from primaite.game.session import PrimaiteSession + +_PRIMAITE_CONFIG["log_level"] = logging.DEBUG print(PRIMAITE_PATHS.app_log_dir_path) import itertools -from primaite.game.session import PrimaiteSession -from primaite.simulator.sim_container import Simulation -from primaite.game.agent.interface import AbstractAgent -from primaite.simulator.network.networks import arcd_uc2_network + import yaml -with open('example_config.yaml', 'r') as file: +from primaite.game.agent.interface import AbstractAgent +from primaite.game.session import PrimaiteSession +from primaite.simulator.network.networks import arcd_uc2_network +from primaite.simulator.sim_container import Simulation + +with open("example_config.yaml", "r") as file: cfg = yaml.safe_load(file) sess = PrimaiteSession.from_config(cfg) -sess.start_session() \ No newline at end of file +sess.start_session() diff --git a/src/primaite/game/agent/GATE_agents.py b/src/primaite/game/agent/GATE_agents.py index eb3c2987..e50d7831 100644 --- a/src/primaite/game/agent/GATE_agents.py +++ b/src/primaite/game/agent/GATE_agents.py @@ -1,9 +1,13 @@ +# flake8: noqa from typing import Dict, Optional, Tuple + +from gymnasium.core import ActType, ObsType + from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractGATEAgent, ObsType from primaite.game.agent.observations import ObservationSpace from primaite.game.agent.rewards import RewardFunction -from gymnasium.core import ActType, ObsType + class GATERLAgent(AbstractGATEAgent): ... @@ -13,9 +17,15 @@ class GATERLAgent(AbstractGATEAgent): # For example MultiAgentEnv in Ray allows sending a dict of observations of multiple agents, then it will reply # with the actions for those agents. - def __init__(self, agent_name: str | None, action_space: ActionManager | None, observation_space: ObservationSpace | None, reward_function: RewardFunction | None) -> None: + def __init__( + self, + agent_name: str | None, + action_space: ActionManager | None, + observation_space: ObservationSpace | None, + reward_function: RewardFunction | None, + ) -> None: super().__init__(agent_name, action_space, observation_space, reward_function) self.most_recent_action: ActType def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: - return self.most_recent_action \ No newline at end of file + return self.most_recent_action diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 4c4aaab4..0a380487 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -695,14 +695,14 @@ class ActionManager: return {i: p for i, p in enumerate(all_action_possibilities)} def get_action(self, action: int) -> Tuple[str, Dict]: - """Produce action in CAOS format""" + """Produce action in CAOS format.""" """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" """The caos format is basically a action identifier, followed by parameters stored in a dictionary""" act_identifier, act_options = self.action_map[action] return act_identifier, act_options - def form_request(self, action_identifier: str, action_options: Dict): - """Take action in CAOS format and use the execution definition to change it into PrimAITE request format""" + def form_request(self, action_identifier: str, action_options: Dict) -> List[str]: + """Take action in CAOS format and use the execution definition to change it into PrimAITE request format.""" act_obj = self.actions[action_identifier] return act_obj.form_request(**action_options) @@ -712,7 +712,8 @@ class ActionManager: return spaces.Discrete(len(self.action_map)) def get_node_uuid_by_idx(self, node_idx: int) -> str: - """Get the node UUID corresponding to the given index. + """ + Get the node UUID corresponding to the given index. :param node_idx: The index of the node to retrieve. :type node_idx: int @@ -722,7 +723,8 @@ class ActionManager: return self.node_uuids[node_idx] def get_folder_uuid_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]: - """Get the folder UUID corresponding to the given node and folder indices. + """ + Get the folder UUID corresponding to the given node and folder indices. :param node_idx: The index of the node. :type node_idx: int @@ -731,7 +733,6 @@ class ActionManager: :return: The UUID of the folder. Or None if the node has fewer folders than the given index. :rtype: Optional[str] """ - node_uuid = self.get_node_uuid_by_idx(node_idx) node = self.sim.network.nodes[node_uuid] folder_uuids = list(node.file_system.folders.keys()) diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 5f121fcc..89f27f3f 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -1,6 +1,6 @@ """Interface for agents.""" from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple, TypeAlias, Union +from typing import Dict, List, Optional, Tuple, TypeAlias, Union import numpy as np @@ -21,6 +21,18 @@ class AbstractAgent(ABC): observation_space: Optional[ObservationSpace], reward_function: Optional[RewardFunction], ) -> None: + """ + Initialize an agent. + + :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes. + :type agent_name: Optional[str] + :param action_space: Action space for the agent. + :type action_space: Optional[ActionManager] + :param observation_space: Observation space for the agent. + :type observation_space: Optional[ObservationSpace] + :param reward_function: Reward function for the agent. + :type reward_function: Optional[RewardFunction] + """ self.agent_name: str = agent_name or "unnamed_agent" self.action_space: Optional[ActionManager] = action_space self.observation_space: Optional[ObservationSpace] = observation_space @@ -32,16 +44,38 @@ class AbstractAgent(ABC): def convert_state_to_obs(self, state: Dict) -> ObsType: """ + Convert a state from the simulator into an observation for the agent using the observation space. + state : dict state directly from simulation.describe_state output : dict state according to CAOS. """ return self.observation_space.observe(state) def calculate_reward_from_state(self, state: Dict) -> float: + """ + Use the reward function to calculate a reward from the state. + + :param state: State of the environment. + :type state: Dict + :return: Reward from the state. + :rtype: float + """ return self.reward_function.calculate(state) @abstractmethod def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: + """ + Return an action to be taken in the environment. + + Subclasses should implement agent logic here. It should use the observation as input to decide best next action. + + :param obs: Observation of the environment. + :type obs: ObsType + :param reward: Reward from the previous action, defaults to None TODO: should this parameter even be accepted? + :type reward: float, optional + :return: Action to be taken in the environment. + :rtype: Tuple[str, Dict] + """ # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 0-39, # then use a bespoke conversion to take 1-40 int back into CAOS action return ("DO_NOTHING", {}) @@ -64,6 +98,15 @@ class RandomAgent(AbstractScriptedAgent): """Agent that ignores its observation and acts completely at random.""" def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: + """Randomly sample an action from the action space. + + :param obs: _description_ + :type obs: ObsType + :param reward: _description_, defaults to None + :type reward: float, optional + :return: _description_ + :rtype: Tuple[str, Dict] + """ return self.action_space.get_action(self.action_space.space.sample()) diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index ba1e8e66..af398fc9 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -1,12 +1,12 @@ +"""Manages the observation space for the agent.""" from abc import ABC, abstractmethod -from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING +from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from gymnasium import spaces -from pydantic import BaseModel -from primaite.simulator.sim_container import Simulation -from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite import getLogger +from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE + _LOGGER = getLogger(__name__) if TYPE_CHECKING: @@ -14,22 +14,25 @@ if TYPE_CHECKING: class AbstractObservation(ABC): + """Abstract class for an observation space component.""" + @abstractmethod def observe(self, state: Dict) -> Any: - """_summary_ + """ + Return an observation based on the current state of the simulation. - :param state: _description_ + :param state: Simulation state dictionary :type state: Dict - :return: _description_ + :return: Observation :rtype: Any """ - ... + pass @property @abstractmethod def space(self) -> spaces.Space: - """Subclasses must define the shape that they expect""" - ... + """Gymnasium space object describing the observation space.""" + pass @classmethod @abstractmethod @@ -39,12 +42,15 @@ class AbstractObservation(ABC): The `session` parameter is for a the PrimaiteSession object that spawns this component. During deserialisation, a subclass of this class may need to translate from a 'reference' to a UUID. """ + pass class FileObservation(AbstractObservation): + """Observation of a file on a node in the network.""" + def __init__(self, where: Optional[Tuple[str]] = None) -> None: """ - _summary_ + Initialise file observation. :param where: Store information about where in the simulation state dictionary to find the relevatn information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with @@ -60,6 +66,13 @@ class FileObservation(AbstractObservation): "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted." def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation file_state = access_from_nested_dict(state, self.where) @@ -69,19 +82,38 @@ class FileObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape. + + :return: Gymnasium space + :rtype: spaces.Space + """ return spaces.Dict({"health_status": spaces.Discrete(6)}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None): + def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: List[str] = None) -> "FileObservation": + """Create file observation from a config. + + :param config: Dictionary containing the configuration for this file observation. + :type config: Dict + :param session: _description_ + :type session: PrimaiteSession + :param parent_where: _description_, defaults to None + :type parent_where: _type_, optional + :return: _description_ + :rtype: _type_ + """ return cls(where=parent_where + ["files", config["file_name"]]) class ServiceObservation(AbstractObservation): + """Observation of a service in the network.""" + default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} "Default observation is what should be returned when the service doesn't exist." def __init__(self, where: Optional[Tuple[str]] = None) -> None: - """ + """Initialise service observation. + :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with zeroes. @@ -94,6 +126,13 @@ class ServiceObservation(AbstractObservation): self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation @@ -104,19 +143,36 @@ class ServiceObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None): + def from_config( + cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None + ) -> "ServiceObservation": + """Create service observation from a config. + + :param config: Dictionary containing the configuration for this service observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional. + :type parent_where: Optional[List[str]], optional + :return: Constructed service observation + :rtype: ServiceObservation + """ return cls(where=parent_where + ["services", session.ref_map_services[config["service_ref"]].uuid]) class LinkObservation(AbstractObservation): + """Observation of a link in the network.""" + default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} "Default observation is what should be returned when the link doesn't exist." def __init__(self, where: Optional[Tuple[str]] = None) -> None: - """ + """Initialise link observation. + :param where: Store information about where in the simulation state dictionary to find the relevant information. Optional. If None, this corresponds that the file does not exist and the observation will be populated with zeroes. @@ -129,6 +185,13 @@ class LinkObservation(AbstractObservation): self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation @@ -147,15 +210,33 @@ class LinkObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape. + + :return: Gymnasium space + :rtype: spaces.Space + """ return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession"): + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "LinkObservation": + """Create link observation from a config. + + :param config: Dictionary containing the configuration for this link observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :return: Constructed link observation + :rtype: LinkObservation + """ return cls(where=["network", "links", session.ref_map_links[config["link_ref"]]]) class FolderObservation(AbstractObservation): - def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder:int=2) -> None: + """Folder observation, including files inside of the folder.""" + + def __init__( + self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder: int = 2 + ) -> None: """Initialise folder Observation, including files inside of the folder. :param where: Where in the simulation state dictionary to find the relevant information for this folder. @@ -179,7 +260,7 @@ class FolderObservation(AbstractObservation): self.files: List[FileObservation] = files while len(self.files) < num_files_per_folder: self.files.append(FileObservation()) - while len(self.files)> num_files_per_folder: + while len(self.files) > num_files_per_folder: truncated_file = self.files.pop() msg = f"Too many files in folde observation. Truncating file {truncated_file}" _LOGGER.warn(msg) @@ -191,6 +272,13 @@ class FolderObservation(AbstractObservation): } def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation folder_state = access_from_nested_dict(state, self.where) @@ -208,6 +296,11 @@ class FolderObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape. + + :return: Gymnasium space + :rtype: spaces.Space + """ return spaces.Dict( { "health_status": spaces.Discrete(6), @@ -216,7 +309,26 @@ class FolderObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]], num_files_per_folder:int=2): + def from_config( + cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]], num_files_per_folder: int = 2 + ) -> "FolderObservation": + """Create folder observation from a config. Also creates child file observations. + + :param config: Dictionary containing the configuration for this folder observation. Includes the name of the + folder and the files inside of it. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :param parent_where: Where in the simulation state dictionary to find the information about this folder's + parent node. A typical location for a node ``where`` can be: + ['network','nodes',,'file_system'] + :type parent_where: Optional[List[str]] + :param num_files_per_folder: How many spaces for files are in this folder observation (to preserve static + observation size) , defaults to 2 + :type num_files_per_folder: int, optional + :return: Constructed folder observation + :rtype: FolderObservation + """ where = parent_where + ["folders", config["folder_name"]] file_configs = config["files"] @@ -226,13 +338,30 @@ class FolderObservation(AbstractObservation): class NicObservation(AbstractObservation): + """Observation of a Network Interface Card (NIC) in the network.""" + default_observation: spaces.Space = {"nic_status": 0} def __init__(self, where: Optional[Tuple[str]] = None) -> None: + """Initialise NIC observation. + + :param where: Where in the simulation state dictionary to find the relevant information for this NIC. A typical + example may look like this: + ['network','nodes',,'NICs',] + If None, this denotes that the NIC does not exist and the observation will be populated with zeroes. + :type where: Optional[Tuple[str]], optional + """ super().__init__() self.where: Optional[Tuple[str]] = where def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation nic_state = access_from_nested_dict(state, self.where) @@ -243,14 +372,31 @@ class NicObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" return spaces.Dict({"nic_status": spaces.Discrete(3)}) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]): + def from_config( + cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] + ) -> "NicObservation": + """Create NIC observation from a config. + + :param config: Dictionary containing the configuration for this NIC observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :param parent_where: Where in the simulation state dictionary to find the information about this NIC's parent + node. A typical location for a node ``where`` can be: ['network','nodes',] + :type parent_where: Optional[List[str]] + :return: Constructed NIC observation + :rtype: NicObservation + """ return cls(where=parent_where + ["NICs", config["nic_uuid"]]) class NodeObservation(AbstractObservation): + """Observation of a node in the network. Includes services, folders and NICs.""" + def __init__( self, where: Optional[Tuple[str]] = None, @@ -260,8 +406,8 @@ class NodeObservation(AbstractObservation): logon_status: bool = False, num_services_per_node: int = 2, num_folders_per_node: int = 2, - num_files_per_folder: int = 2 - ) -> None: + num_files_per_folder: int = 2, + ) -> None: """ Configurable observation for a node in the simulation. @@ -271,7 +417,8 @@ class NodeObservation(AbstractObservation): :type where: List[str], optional :param services: Mapping between position in observation space and service UUID, defaults to {} :type services: Dict[int,str], optional - :param max_services: Max number of services that can be presented in observation space for this node, defaults to 2 + :param max_services: Max number of services that can be presented in observation space for this node + , defaults to 2 :type max_services: int, optional :param folders: Mapping between position in observation space and folder name, defaults to {} :type folders: Dict[int,str], optional @@ -286,10 +433,10 @@ class NodeObservation(AbstractObservation): self.where: Optional[Tuple[str]] = where self.services: List[ServiceObservation] = services - while len(self.services)num_services_per_node: + while len(self.services) > num_services_per_node: truncated_service = self.services.pop() msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}" _LOGGER.warn(msg) @@ -297,8 +444,8 @@ class NodeObservation(AbstractObservation): # truncate service list self.folders: List[FolderObservation] = folders + # add empty folder observation without `where` parameter that will always return default (blank) observations while len(self.folders) < num_folders_per_node: - # add an empty folder observation without `where` parameter that will always return default (blank) observations self.folders.append(FolderObservation()) while len(self.folders) > num_folders_per_node: truncated_folder = self.folders.pop() @@ -317,6 +464,13 @@ class NodeObservation(AbstractObservation): self.default_observation["logon_status"] = 0 def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation @@ -337,6 +491,7 @@ class NodeObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" space_shape = { "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}), "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}), @@ -358,6 +513,27 @@ class NodeObservation(AbstractObservation): num_folders_per_node: int = 2, num_files_per_folder: int = 2, ) -> "NodeObservation": + """Create node observation from a config. Also creates child service, folder and NIC observations. + + :param config: Dictionary containing the configuration for this node observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :param parent_where: Where in the simulation state dictionary to find the information about this node's parent + network. A typical location for it would be: ['network',] + :type parent_where: Optional[List[str]] + :param num_services_per_node: How many spaces for services are in this node observation (to preserve static + observation size) , defaults to 2 + :type num_services_per_node: int, optional + :param num_folders_per_node: How many spaces for folders are in this node observation (to preserve static + observation size) , defaults to 2 + :type num_folders_per_node: int, optional + :param num_files_per_folder: How many spaces for files are in the folder observations (to preserve static + observation size) , defaults to 2 + :type num_files_per_folder: int, optional + :return: Constructed node observation + :rtype: NodeObservation + """ node_uuid = session.ref_map_nodes[config["node_ref"]] if parent_where is None: where = ["network", "nodes", node_uuid] @@ -367,7 +543,12 @@ class NodeObservation(AbstractObservation): svc_configs = config.get("services", {}) services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs] folder_configs = config.get("folders", {}) - folders = [FolderObservation.from_config(config=c, session=session, parent_where=where, num_files_per_folder=num_files_per_folder) for c in folder_configs] + folders = [ + FolderObservation.from_config( + config=c, session=session, parent_where=where, num_files_per_folder=num_files_per_folder + ) + for c in folder_configs + ] nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys() nic_configs = [{"nic_uuid": n for n in nic_uuids}] if nic_uuids else [] nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs] @@ -378,13 +559,15 @@ class NodeObservation(AbstractObservation): folders=folders, nics=nics, logon_status=logon_status, - num_services_per_node = num_services_per_node, - num_folders_per_node = num_folders_per_node, - num_files_per_folder = num_files_per_folder, - ) + num_services_per_node=num_services_per_node, + num_folders_per_node=num_folders_per_node, + num_files_per_folder=num_files_per_folder, + ) class AclObservation(AbstractObservation): + """Observation of an Access Control List (ACL) in the network.""" + # TODO: should where be optional, and we can use where=None to pad the observation space? # definitely the current approach does not support tracking files that aren't specified by name, for example # if a file is created at runtime, we have currently got no way of telling the observation space to track it. @@ -397,6 +580,21 @@ class AclObservation(AbstractObservation): where: Optional[Tuple[str]] = None, num_rules: int = 10, ) -> None: + """Initialise ACL observation. + + :param node_ip_to_id: Mapping between IP address and ID. + :type node_ip_to_id: Dict[str, int] + :param ports: List of ports which are part of the game that define the ordering when converting to an ID + :type ports: List[int] + :param protocols: List of protocols which are part of the game, defines ordering when converting to an ID + :type protocols: list[str] + :param where: Where in the simulation state dictionary to find the relevant information for this ACL. A typical + example may look like this: + ['network','nodes',,'acl','acl'] + :type where: Optional[Tuple[str]], optional + :param num_rules: , defaults to 10 + :type num_rules: int, optional + """ super().__init__() self.where: Optional[Tuple[str]] = where self.num_rules: int = num_rules @@ -423,6 +621,13 @@ class AclObservation(AbstractObservation): } def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation acl_state: Dict = access_from_nested_dict(state, self.where) @@ -457,6 +662,11 @@ class AclObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape. + + :return: Gymnasium space + :rtype: spaces.Space + """ return spaces.Dict( { "RULES": spaces.Dict( @@ -482,6 +692,15 @@ class AclObservation(AbstractObservation): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation": + """Generate ACL observation from a config. + + :param config: Dictionary containing the configuration for this ACL observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :return: Observation object + :rtype: AclObservation + """ node_ip_to_idx = {} for ip_idx, ip_map_config in enumerate(config["ip_address_order"]): node_ref = ip_map_config["node_ref"] @@ -500,26 +719,44 @@ class AclObservation(AbstractObservation): class NullObservation(AbstractObservation): + """Null observation, returns a single 0 value for the observation space.""" + def __init__(self, where: Optional[List[str]] = None): + """Initialise null observation.""" self.default_observation: Dict = {} def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation.""" return 0 @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" return spaces.Discrete(1) @classmethod def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation": + """ + Create null observation from a config. + + The parameters are ignored, they are here to match the signature of the other observation classes. + """ return cls() class ICSObservation(NullObservation): + """ICS observation placeholder, currently not implemented so always returns a single 0.""" + pass class UC2BlueObservation(AbstractObservation): + """Container for all observations used by the blue agent in UC2. + + TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler + for the purpose of compiling several observation components. + """ + def __init__( self, nodes: List[NodeObservation], @@ -528,6 +765,20 @@ class UC2BlueObservation(AbstractObservation): ics: ICSObservation, where: Optional[List[str]] = None, ) -> None: + """Initialise UC2 blue observation. + + :param nodes: List of node observations + :type nodes: List[NodeObservation] + :param links: List of link observations + :type links: List[LinkObservation] + :param acl: The Access Control List observation + :type acl: AclObservation + :param ics: The ICS observation + :type ics: ICSObservation + :param where: Where in the simulation state dict to find information. Not used in this particular observation + because it only compiles other observations and doesn't contribute any new information, defaults to None + :type where: Optional[List[str]], optional + """ super().__init__() self.where: Optional[Tuple[str]] = where @@ -544,6 +795,13 @@ class UC2BlueObservation(AbstractObservation): } def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ if self.where is None: return self.default_observation @@ -557,6 +815,12 @@ class UC2BlueObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """ + Gymnasium space object describing the observation space shape. + + :return: Space + :rtype: spaces.Space + """ return spaces.Dict( { "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), @@ -567,21 +831,34 @@ class UC2BlueObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession"): + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "UC2BlueObservation": + """Create UC2 blue observation from a config. + + :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes, + links, ACL and ICS observations. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + :return: Constructed UC2 blue observation + :rtype: UC2BlueObservation + """ node_configs = config["nodes"] num_services_per_node = config["num_services_per_node"] num_folders_per_node = config["num_folders_per_node"] num_files_per_folder = config["num_files_per_folder"] - nodes = [NodeObservation.from_config( - config=n, - session=session, - num_services_per_node= num_services_per_node, - num_folders_per_node=num_folders_per_node, - num_files_per_folder=num_files_per_folder, - ) for n in node_configs] + nodes = [ + NodeObservation.from_config( + config=n, + session=session, + num_services_per_node=num_services_per_node, + num_folders_per_node=num_folders_per_node, + num_files_per_folder=num_files_per_folder, + ) + for n in node_configs + ] link_configs = config["links"] - links = [LinkObservation.from_config(config=l, session=session) for l in link_configs] + links = [LinkObservation.from_config(config=link, session=session) for link in link_configs] acl_config = config["acl"] acl = AclObservation.from_config(config=acl_config, session=session) @@ -593,6 +870,8 @@ class UC2BlueObservation(AbstractObservation): class UC2RedObservation(AbstractObservation): + """Container for all observations used by the red agent in UC2.""" + def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None: super().__init__() self.where: Optional[List[str]] = where @@ -603,6 +882,7 @@ class UC2RedObservation(AbstractObservation): } def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation.""" if self.where is None: return self.default_observation @@ -612,6 +892,7 @@ class UC2RedObservation(AbstractObservation): @property def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" return spaces.Dict( { "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), @@ -619,43 +900,74 @@ class UC2RedObservation(AbstractObservation): ) @classmethod - def from_config(cls, config: Dict, session: "PrimaiteSession"): + def from_config(cls, config: Dict, session: "PrimaiteSession") -> "UC2RedObservation": + """ + Create UC2 red observation from a config. + + :param config: Dictionary containing the configuration for this UC2 red observation. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + """ node_configs = config["nodes"] nodes = [NodeObservation.from_config(config=cfg, session=session) for cfg in node_configs] return cls(nodes=nodes, where=["network"]) class UC2GreenObservation(NullObservation): + """Green agent observation. As the green agent's actions don't depend on the observation, this is empty.""" + pass class ObservationSpace: """ - Manage the observations of an Actor. + Manage the observations of an Agent. The observation space has the purpose of: 1. Reading the outputted state from the PrimAITE Simulation. 2. Selecting parts of the simulation state that are requested by the simulation config - 3. Formatting this information so an actor can use it to make decisions. + 3. Formatting this information so an agent can use it to make decisions. """ - ... + # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed + # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next + # refactor. - # what this class does: - # keep a list of observations - # create observations for an actor from the config def __init__(self, observation: AbstractObservation) -> None: + """Initialise observation space. + + :param observation: Observation object + :type observation: AbstractObservation + """ self.obs: AbstractObservation = observation - def observe(self, state) -> Dict: + def observe(self, state: Dict) -> Dict: + """ + Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + """ return self.obs.observe(state) @property def space(self) -> None: + """Gymnasium space object describing the observation space shape.""" return self.obs.space @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "ObservationSpace": + """Create observation space from a config. + + :param config: Dictionary containing the configuration for this observation space. + It should contain the key 'type' which selects which observation class to use (from a choice of: + UC2BlueObservation, UC2RedObservation, UC2GreenObservation) + The other key is 'options' which are passed to the constructor of the selected observation class. + :type config: Dict + :param session: Reference to the PrimaiteSession object that spawned this observation. + :type session: PrimaiteSession + """ if config["type"] == "UC2BlueObservation": return cls(UC2BlueObservation.from_config(config.get("options", {}), session=session)) elif config["type"] == "UC2RedObservation": diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 67e6ee50..03c4e2d3 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -1,8 +1,35 @@ -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Tuple, TYPE_CHECKING +""" +Manages the reward function for the agent. + +Each agent is equipped with a RewardFunction, which is made up of a list of reward components. The components are +designed to calculate a reward value based on the current state of the simulation. The overall reward function is a +weighed sum of the components. + +The reward function is typically specified using a config yaml file or a config dictionary. The following example shows +the structure: +```yaml + reward_function: + reward_components: + - type: DATABASE_FILE_INTEGRITY + weight: 0.5 + options: + node_ref: database_server + folder_name: database + file_name: database.db + + + - type: WEB_SERVER_404_PENALTY + weight: 0.5 + options: + node_ref: web_server + service_ref: web_server_database_client +``` +""" +from abc import abstractmethod +from typing import Dict, List, Tuple, TYPE_CHECKING from primaite import getLogger -from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE +from primaite.game.agent.utils import access_from_nested_dict _LOGGER = getLogger(__name__) @@ -11,27 +38,60 @@ if TYPE_CHECKING: class AbstractReward: + """Base class for reward function components.""" + @abstractmethod def calculate(self, state: Dict) -> float: + """Calculate the reward for the current state.""" return 0.0 @classmethod @abstractmethod def from_config(cls, config: dict, session: "PrimaiteSession") -> "AbstractReward": + """Create a reward function component from a config dictionary. + + :param config: dict of options for the reward component's constructor + :type config: dict + :param session: Reference to the PrimAITE Session object + :type session: PrimaiteSession + :return: The reward component. + :rtype: AbstractReward + """ return cls() class DummyReward(AbstractReward): + """Dummy reward function component which always returns 0.""" + def calculate(self, state: Dict) -> float: + """Calculate the reward for the current state.""" return 0.0 @classmethod def from_config(cls, config: dict, session: "PrimaiteSession") -> "DummyReward": + """Create a reward function component from a config dictionary. + + :param config: dict of options for the reward component's constructor. Should be empty. + :type config: dict + :param session: Reference to the PrimAITE Session object + :type session: PrimaiteSession + """ return cls() class DatabaseFileIntegrity(AbstractReward): + """Reward function component which rewards the agent for maintaining the integrity of a database file.""" + def __init__(self, node_uuid: str, folder_name: str, file_name: str) -> None: + """Initialise the reward component. + + :param node_uuid: UUID of the node which contains the database file. + :type node_uuid: str + :param folder_name: folder which contains the database file. + :type folder_name: str + :param file_name: name of the database file. + :type file_name: str + """ self.location_in_state = [ "network", "nodes", @@ -44,6 +104,11 @@ class DatabaseFileIntegrity(AbstractReward): ] def calculate(self, state: Dict) -> float: + """Calculate the reward for the current state. + + :param state: The current state of the simulation. + :type state: Dict + """ database_file_state = access_from_nested_dict(state, self.location_in_state) health_status = database_file_state["health_status"] if health_status == "corrupted": @@ -55,6 +120,15 @@ class DatabaseFileIntegrity(AbstractReward): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "DatabaseFileIntegrity": + """Create a reward function component from a config dictionary. + + :param config: dict of options for the reward component's constructor + :type config: Dict + :param session: Reference to the PrimAITE Session object + :type session: PrimaiteSession + :return: The reward component. + :rtype: DatabaseFileIntegrity + """ node_ref = config.get("node_ref") folder_name = config.get("folder_name") file_name = config.get("file_name") @@ -76,7 +150,10 @@ class DatabaseFileIntegrity(AbstractReward): node_uuid = session.ref_map_nodes[node_ref] if not node_uuid: _LOGGER.error( - f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation" + ( + f"{cls.__name__} could not be initialised from config because the referenced node could not be " + f"found in the simulation" + ) ) return DummyReward() # TODO: better error handling @@ -84,10 +161,24 @@ class DatabaseFileIntegrity(AbstractReward): class WebServer404Penalty(AbstractReward): + """Reward function component which penalises the agent when the web server returns a 404 error.""" + def __init__(self, node_uuid: str, service_uuid: str) -> None: + """Initialise the reward component. + + :param node_uuid: UUID of the node which contains the web server service. + :type node_uuid: str + :param service_uuid: UUID of the web server service. + :type service_uuid: str + """ self.location_in_state = ["network", "nodes", node_uuid, "services", service_uuid] def calculate(self, state: Dict) -> float: + """Calculate the reward for the current state. + + :param state: The current state of the simulation. + :type state: Dict + """ web_service_state = access_from_nested_dict(state, self.location_in_state) most_recent_return_code = web_service_state["most_recent_return_code"] # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan. @@ -100,16 +191,31 @@ class WebServer404Penalty(AbstractReward): @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "WebServer404Penalty": + """Create a reward function component from a config dictionary. + + :param config: dict of options for the reward component's constructor + :type config: Dict + :param session: Reference to the PrimAITE Session object + :type session: PrimaiteSession + :return: The reward component. + :rtype: WebServer404Penalty + """ node_ref = config.get("node_ref") service_ref = config.get("service_ref") if not (node_ref and service_ref): - msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config." + msg = ( + f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not " + "found in reward config." + ) _LOGGER.warn(msg) return DummyReward() # TODO: should we error out with incorrect inputs? Probably! node_uuid = session.ref_map_nodes[node_ref] service_uuid = session.ref_map_services[service_ref].uuid if not (node_uuid and service_uuid): - msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator." + msg = ( + f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not" + " found in the simulator." + ) _LOGGER.warn(msg) return DummyReward() # TODO: consider erroring here as well @@ -117,6 +223,8 @@ class WebServer404Penalty(AbstractReward): class RewardFunction: + """Manages the reward function for the agent.""" + __rew_class_identifiers: Dict[str, type[AbstractReward]] = { "DUMMY": DummyReward, "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity, @@ -124,13 +232,26 @@ class RewardFunction: } def __init__(self): + """Initialise the reward function object.""" self.reward_components: List[Tuple[AbstractReward, float]] = [] "attribute reward_components keeps track of reward components and the weights assigned to each." def regsiter_component(self, component: AbstractReward, weight: float = 1.0) -> None: + """Add a reward component to the reward function. + + :param component: Instance of a reward component. + :type component: AbstractReward + :param weight: Relative weight of the reward component, defaults to 1.0 + :type weight: float, optional + """ self.reward_components.append((component, weight)) def calculate(self, state: Dict) -> float: + """Calculate the overall reward for the current state. + + :param state: The current state of the simulation. + :type state: Dict + """ total = 0.0 for comp_and_weight in self.reward_components: comp = comp_and_weight[0] @@ -140,6 +261,15 @@ class RewardFunction: @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "RewardFunction": + """Create a reward function from a config dictionary. + + :param config: dict of options for the reward manager's constructor + :type config: Dict + :param session: Reference to the PrimAITE Session object + :type session: PrimaiteSession + :return: The reward manager. + :rtype: RewardFunction + """ new = cls() for rew_component_cfg in config["reward_components"]: diff --git a/src/primaite/game/agent/scripted_agents.py b/src/primaite/game/agent/scripted_agents.py index d3becd57..3748494b 100644 --- a/src/primaite/game/agent/scripted_agents.py +++ b/src/primaite/game/agent/scripted_agents.py @@ -1,9 +1,14 @@ +"""Agents with predefined behaviours.""" from primaite.game.agent.interface import AbstractScriptedAgent class GreenWebBrowsingAgent(AbstractScriptedAgent): - ... + """Scripted agent which attempts to send web requests to a target node.""" + + raise NotImplementedError class RedDatabaseCorruptingAgent(AbstractScriptedAgent): - ... + """Scripted agent which attempts to corrupt the database of the target node.""" + + raise NotImplementedError diff --git a/src/primaite/game/agent/utils.py b/src/primaite/game/agent/utils.py index ad6dbefe..1314087c 100644 --- a/src/primaite/game/agent/utils.py +++ b/src/primaite/game/agent/utils.py @@ -1,4 +1,4 @@ -from typing import Dict, Sequence, Hashable, Any +from typing import Any, Dict, Hashable, Sequence NOT_PRESENT_IN_STATE = object() """ @@ -6,6 +6,7 @@ Need an object to return when the sim state does not contain a requested value. the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is a sentinel for this purpose. """ + def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: """ Access an item from a deeply dictionary with a list of keys. @@ -26,4 +27,4 @@ def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any: k = key_list.pop(0) if k not in dictionary: return NOT_PRESENT_IN_STATE - return access_from_nested_dict(dictionary[k], key_list) \ No newline at end of file + return access_from_nested_dict(dictionary[k], key_list) diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index f29d03dd..bd5e18d6 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -1,6 +1,6 @@ """PrimAITE session - the main entry point to training agents on PrimAITE.""" from ipaddress import IPv4Address -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Never, Optional, Tuple from arcd_gate.client.gate_client import ActType, GATEClient from gymnasium import spaces @@ -32,13 +32,17 @@ _LOGGER = getLogger(__name__) class PrimaiteGATEClient(GATEClient): + """Lightweight wrapper around the GATEClient class that allows PrimAITE to message GATE.""" + def __init__(self, parent_session: "PrimaiteSession", service_port: int = 50000): - """Create a new GATE client for PrimAITE. + """ + Create a new GATE client for PrimAITE. :param parent_session: The parent session object. :type parent_session: PrimaiteSession :param service_port: The port on which the GATE service is running. - :type service_port: int, optional""" + :type service_port: int, optional + """ super().__init__(service_port=service_port) self.parent_session: "PrimaiteSession" = parent_session @@ -133,9 +137,11 @@ class PrimaiteGATEClient(GATEClient): class PrimaiteSessionOptions(BaseModel): - """Global options which are applicable to all of the agents in the game. + """ + Global options which are applicable to all of the agents in the game. - Currently this is used to restrict which ports and protocols exist in the world of the simulation.""" + Currently this is used to restrict which ports and protocols exist in the world of the simulation. + """ ports: List[str] protocols: List[str] @@ -154,9 +160,7 @@ class TrainingOptions(BaseModel): class PrimaiteSession: - """ - The main entrypoint for PrimAITE sessions, this coordinates a simulation, agents, and connections to ARCD GATE. - """ + """The main entrypoint for PrimAITE sessions, this manages a simulation, agents, and connections to ARCD GATE.""" def __init__(self): self.simulation: Simulation = Simulation() @@ -183,7 +187,7 @@ class PrimaiteSession: self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self) """Reference to a GATE Client object, which will send data to GATE service for training RL agent.""" - def start_session(self, opts="TODO..."): + def start_session(self) -> Never: """Commence the training session, this gives the GATE client control over the simulation/agent loop.""" self.gate_client.start() @@ -221,7 +225,7 @@ class PrimaiteSession: # to discrete(40) is only necessary for purposes of RL learning, therefore that bit of # code should live inside of the GATE agent subclass) # gets action in CAOS format - _LOGGER.debug(f"Getting agent action") + _LOGGER.debug("Getting agent action") agent_action, action_options = agent.get_action(agent_obs, agent_reward) # 9. CAOS action is converted into request (extra information might be needed to enrich # the request, this is what the execution definition is there for) @@ -236,11 +240,11 @@ class PrimaiteSession: self.simulation.apply_timestep(self.step_counter) self.step_counter += 1 - def reset(self): + def reset(self) -> None: """Reset the session, this will reset the simulation.""" return NotImplemented - def close(self): + def close(self) -> None: """Close the session, this will stop the gate client and close the simulation.""" return NotImplemented From 0f24b4a646ae65bc7b1e4afc230133a759dadd06 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 19 Oct 2023 15:34:46 +0100 Subject: [PATCH 26/35] Remove broken import --- src/primaite/game/session.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index bd5e18d6..adb9f7b5 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -1,6 +1,6 @@ """PrimAITE session - the main entry point to training agents on PrimAITE.""" from ipaddress import IPv4Address -from typing import Any, Dict, List, Never, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from arcd_gate.client.gate_client import ActType, GATEClient from gymnasium import spaces @@ -187,7 +187,7 @@ class PrimaiteSession: self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self) """Reference to a GATE Client object, which will send data to GATE service for training RL agent.""" - def start_session(self) -> Never: + def start_session(self) -> None: """Commence the training session, this gives the GATE client control over the simulation/agent loop.""" self.gate_client.start() From 975aa9ffc289271ead984b4b94b06adf7d2011d6 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 23 Oct 2023 16:26:34 +0100 Subject: [PATCH 27/35] Minor changes to rewards and services. --- example_config.yaml | 14 +++++++----- src/primaite/game/agent/rewards.py | 13 ++++++----- src/primaite/game/session.py | 22 +++++++++++++++++-- .../system/applications/web_browser.py | 3 ++- 4 files changed, 39 insertions(+), 13 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index e16411fa..f3d8dc10 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -2,10 +2,10 @@ training_config: rl_framework: SB3 rl_algorithm: PPO seed: 333 - n_learn_episodes: 1 - n_learn_steps: 8 - n_eval_episodes: 0 - n_eval_steps: 8 + n_learn_episodes: 4 + n_learn_steps: 128 + n_eval_episodes: 1 + n_eval_steps: 128 game_config: @@ -534,6 +534,9 @@ simulation: type: DatabaseClient options: db_server_ip: 192.168.1.14 + - ref: web_server_web_service + type: WebServer + - ref: database_server type: server @@ -589,9 +592,10 @@ simulation: subnet_mask: 255.255.255.0 default_gateway: 192.168.10.1 dns_server: 192.168.1.10 - services: + applications: - ref: client_2_web_browser type: WebBrowser + services: - ref: client_2_dns_client type: DNSClient diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 03c4e2d3..6c408ff9 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -29,7 +29,7 @@ from abc import abstractmethod from typing import Dict, List, Tuple, TYPE_CHECKING from primaite import getLogger -from primaite.game.agent.utils import access_from_nested_dict +from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE _LOGGER = getLogger(__name__) @@ -180,14 +180,17 @@ class WebServer404Penalty(AbstractReward): :type state: Dict """ web_service_state = access_from_nested_dict(state, self.location_in_state) - most_recent_return_code = web_service_state["most_recent_return_code"] + if web_service_state is NOT_PRESENT_IN_STATE: + print("error getting web service state") + return 0.0 + most_recent_return_code = web_service_state["last_response_status_code"] # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan. if most_recent_return_code == 200: - return 1 + return 1.0 elif most_recent_return_code == 404: - return -1 + return -1.0 else: - return 0 + return 0.0 @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "WebServer404Penalty": diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index adb9f7b5..d40d0754 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -21,12 +21,15 @@ from primaite.simulator.network.hardware.nodes.switch import Switch from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.sim_container import Simulation +from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.database_client import DatabaseClient +from primaite.simulator.system.applications.web_browser import WebBrowser from primaite.simulator.system.services.database.database_service import DatabaseService from primaite.simulator.system.services.dns.dns_client import DNSClient from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.service import Service +from primaite.simulator.system.services.web_server.web_server import WebServer _LOGGER = getLogger(__name__) @@ -182,6 +185,8 @@ class PrimaiteSession: """Mapping from unique node reference name to node object. Used when parsing config files.""" self.ref_map_services: Dict[str, Service] = {} """Mapping from human-readable service reference to service object. Used for parsing config files.""" + self.ref_map_applications: Dict[str, Application] = {} + """Mapping from human-readable application reference to application object. Used for parsing config files.""" self.ref_map_links: Dict[str, Link] = {} """Mapping from human-readable link reference to link object. Used when parsing config files.""" self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self) @@ -333,11 +338,11 @@ class PrimaiteSession: "DNSServer": DNSServer, "DatabaseClient": DatabaseClient, "DatabaseService": DatabaseService, - # 'database_backup': , + "WebServer": WebServer, "DataManipulationBot": DataManipulationBot, - # 'web_browser' } if service_type in service_types_mapping: + print(f"installing {service_type} on node {new_node.hostname}") new_node.software_manager.install(service_types_mapping[service_type]) new_service = new_node.software_manager.software[service_type] sess.ref_map_services[service_ref] = new_service @@ -355,6 +360,19 @@ class PrimaiteSession: if "domain_mapping" in opt: for domain, ip in opt["domain_mapping"].items(): new_service.dns_register(domain, ip) + if "applications" in node_cfg: + for application_cfg in node_cfg["applications"]: + application_ref = application_cfg["ref"] + application_type = application_cfg["type"] + application_types_mapping = { + "WebBrowser": WebBrowser, + } + if application_type in application_types_mapping: + new_node.software_manager.install(application_types_mapping[application_type]) + new_application = new_node.software_manager.software[application_type] + sess.ref_map_applications[application_ref] = new_application + else: + print(f"application type not found {application_type}") if "nics" in node_cfg: for nic_num, nic_cfg in node_cfg["nics"].items(): new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"])) diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index c48b785e..ea9c3ac3 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -38,7 +38,8 @@ class WebBrowser(Application): :return: A dictionary capturing the current state of the WebBrowser and its child objects. """ - return super().describe_state() + state = super().describe_state() + state["last_response_status_code"] = self.latest_response.status_code if self.latest_response else None def reset_component_for_episode(self, episode: int): """ From d4eee36b7bc921af0791e9b1a4aa062a0b97111d Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 23 Oct 2023 17:23:14 +0100 Subject: [PATCH 28/35] Fix software registration for game layer and simulator interface --- example_config.yaml | 2 +- src/primaite/game/agent/observations.py | 2 +- .../simulator/network/hardware/base.py | 36 ++++++++++++++++++- .../system/applications/application.py | 2 +- .../simulator/system/core/software_manager.py | 15 ++++++++ .../system/services/web_server/web_server.py | 20 ++++++++++- src/primaite/simulator/system/software.py | 6 ++-- 7 files changed, 75 insertions(+), 8 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index f3d8dc10..00beaa1e 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -457,7 +457,7 @@ game_config: weight: 0.5 options: node_ref: web_server - service_ref: web_server_database_client + service_ref: web_server_web_service agent_settings: diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index af398fc9..35fe8ac5 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -139,7 +139,7 @@ class ServiceObservation(AbstractObservation): service_state = access_from_nested_dict(state, self.where) if service_state is NOT_PRESENT_IN_STATE: return self.default_observation - return {"operating_status": service_state["operating_state"], "health_status": service_state["health_status"]} + return {"operating_status": service_state["operating_state"], "health_status": service_state["health_state"]} @property def space(self) -> spaces.Space: diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 78ae228e..607e348b 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -938,6 +938,7 @@ class Node(SimComponent): kwargs["file_system"] = FileSystem(sys_log=kwargs["sys_log"], sim_root=kwargs["root"] / "fs") if not kwargs.get("software_manager"): kwargs["software_manager"] = SoftwareManager( + parent_node=self, sys_log=kwargs.get("sys_log"), session_manager=kwargs.get("session_manager"), file_system=kwargs.get("file_system"), @@ -1199,7 +1200,8 @@ class Node(SimComponent): self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager)) def uninstall_service(self, service: Service) -> None: - """Uninstall and completely remove service from this node. + """ + Uninstall and completely remove service from this node. :param service: Service object that is currently associated with this node. :type service: Service @@ -1214,6 +1216,38 @@ class Node(SimComponent): _LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}") self._service_request_manager.remove_request(service.uuid) + def install_application(self, application: Application) -> None: + """ + Install an application on this node. + + :param application: Application instance that has not been installed on any node yet. + :type application: Application + """ + if application in self: + _LOGGER.warning(f"Can't add application {application.uuid} to node {self.uuid}. It's already installed.") + return + self.applications[application.uuid] = application + application.parent = self + self.sys_log.info(f"Installed application {application.name}") + _LOGGER.info(f"Added application {application.uuid} to node {self.uuid}") + self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager)) + + def uninstall_application(self, application: Application) -> None: + """ + Uninstall and completely remove application from this node. + + :param application: Application object that is currently associated with this node. + :type application: Application + """ + if application not in self: + _LOGGER.warning(f"Can't remove application {application.uuid} from node {self.uuid}. It's not installed.") + return + self.applications.pop(application.uuid) + application.parent = None + self.sys_log.info(f"Uninstalled application {application.name}") + _LOGGER.info(f"Removed application {application.uuid} from node {self.uuid}") + self._application_request_manager.remove_request(application.uuid) + def __contains__(self, item: Any) -> bool: if isinstance(item, Service): return item.uuid in self.services diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 69b64aac..e3da6f01 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -45,7 +45,7 @@ class Application(IOSoftware): state = super().describe_state() state.update( { - "opearting_state": self.operating_state.name, + "opearting_state": self.operating_state.value, "execution_control_status": self.execution_control_status, "num_executions": self.num_executions, "groups": list(self.groups), diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 973b17b4..8b8fe599 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -14,6 +14,7 @@ from primaite.simulator.system.software import IOSoftware if TYPE_CHECKING: from primaite.simulator.system.core.session_manager import SessionManager from primaite.simulator.system.core.sys_log import SysLog + from primaite.simulator.network.hardware.base import Node from typing import Type, TypeVar @@ -25,6 +26,7 @@ class SoftwareManager: def __init__( self, + parent_node: "Node", session_manager: "SessionManager", sys_log: SysLog, file_system: FileSystem, @@ -35,6 +37,7 @@ class SoftwareManager: :param session_manager: The session manager handling network communications. """ + self.node = parent_node self.session_manager = session_manager self.software: Dict[str, Union[Service, Application]] = {} self._software_class_to_name_map: Dict[Type[IOSoftwareClass], str] = {} @@ -62,6 +65,8 @@ class SoftwareManager: :param software_class: The software class. """ + # TODO: Software manager and node itself both have an install method. Need to refactor to have more logical + # separation of concerns. if software_class in self._software_class_to_name_map: self.sys_log.info(f"Cannot install {software_class} as it is already installed") return @@ -77,6 +82,12 @@ class SoftwareManager: if isinstance(software, Application): software.operating_state = ApplicationOperatingState.CLOSED + # add the software to the node's registry after it has been fully initialized + if isinstance(software, Service): + self.node.install_service(software) + elif isinstance(software, Application): + self.node.install_application(software) + def uninstall(self, software_name: str): """ Uninstall an Application or Service. @@ -85,6 +96,10 @@ class SoftwareManager: """ if software_name in self.software: software = self.software.pop(software_name) # noqa + if isinstance(software, Application): + self.node.uninstall_application(software) + elif isinstance(software, Service): + self.node.uninstall_service(software) del software self.sys_log.info(f"Deleted {software_name}") return diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py index f63d5169..5957e4cb 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/src/primaite/simulator/system/services/web_server/web_server.py @@ -1,5 +1,5 @@ from ipaddress import IPv4Address -from typing import Any, Optional +from typing import Any, Dict, Optional from urllib.parse import urlparse from primaite.simulator.network.protocols.http import ( @@ -17,6 +17,23 @@ from primaite.simulator.system.services.service import Service class WebServer(Service): """Class used to represent a Web Server Service in simulation.""" + last_response_status_code: Optional[HttpStatusCode] = None + + def describe_state(self) -> Dict: + """ + Produce a dictionary describing the current state of this object. + + Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation. + + :return: Current state of this object and child objects. + :rtype: Dict + """ + state = super().describe_state() + state["last_response_status_code"] = ( + self.last_response_status_code.value if self.last_response_status_code else None + ) + return state + def __init__(self, **kwargs): kwargs["name"] = "WebServer" kwargs["protocol"] = IPProtocol.TCP @@ -66,6 +83,7 @@ class WebServer(Service): self.send(payload=response, session_id=session_id) # return true if response is OK + self.last_response_status_code = response.status_code return response.status_code == HttpStatusCode.OK def _handle_get_request(self, payload: HttpRequestPacket) -> HttpResponsePacket: diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 25f764e4..8cd13d1a 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -119,9 +119,9 @@ class Software(SimComponent): state = super().describe_state() state.update( { - "health_state": self.health_state_actual.name, - "health_state_red_view": self.health_state_visible.name, - "criticality": self.criticality.name, + "health_state": self.health_state_actual.value, + "health_state_red_view": self.health_state_visible.value, + "criticality": self.criticality.value, "patching_count": self.patching_count, "scanning_count": self.scanning_count, "revealed_to_red": self.revealed_to_red, From 6b7c483a678f1941f3fa8998f13e010352d04f16 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 24 Oct 2023 11:07:25 +0100 Subject: [PATCH 29/35] Align observations to Common approach --- example_config.yaml | 31 +++++----- src/primaite/game/agent/observations.py | 78 +++++++++++++------------ 2 files changed, 59 insertions(+), 50 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index 00beaa1e..bcc819ae 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -2,10 +2,10 @@ training_config: rl_framework: SB3 rl_algorithm: PPO seed: 333 - n_learn_episodes: 4 - n_learn_steps: 128 - n_eval_episodes: 1 - n_eval_steps: 128 + n_learn_episodes: 1 + n_learn_steps: 8 + n_eval_episodes: 0 + n_eval_steps: 8 game_config: @@ -39,10 +39,10 @@ game_config: options: nodes: - node_ref: client_2 - max_folders_per_node: 2 - max_files_per_folder: 2 - max_services_per_node: 2 - max_nics_per_node: 8 + max_folders_per_node: 1 + max_files_per_folder: 1 + max_services_per_node: 1 + max_nics_per_node: 2 max_acl_rules: 10 reward_function: @@ -93,9 +93,9 @@ game_config: options: nodes: - node_ref: client_1 - max_folders_per_node: 2 - max_files_per_folder: 2 - max_services_per_node: 2 + max_folders_per_node: 1 + max_files_per_folder: 1 + max_services_per_node: 1 reward_function: reward_components: @@ -113,9 +113,10 @@ game_config: observation_space: type: UC2BlueObservation options: - num_services_per_node: 2 - num_folders_per_node: 2 - num_files_per_folder: 2 + num_services_per_node: 1 + num_folders_per_node: 1 + num_files_per_folder: 1 + num_nics_per_node: 2 nodes: - node_ref: domain_controller services: @@ -148,6 +149,8 @@ game_config: - link_ref: switch_2___client_2 - link_ref: switch_2___security_suite acl: + options: + max_acl_rules: 10 router_node_ref: router_1 ip_address_order: - node_ref: domain_controller diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 35fe8ac5..8eb322bd 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -167,7 +167,7 @@ class ServiceObservation(AbstractObservation): class LinkObservation(AbstractObservation): """Observation of a link in the network.""" - default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}} + default_observation: spaces.Space = {"PROTOCOLS": {"ALL": 0}} "Default observation is what should be returned when the link doesn't exist." def __init__(self, where: Optional[Tuple[str]] = None) -> None: @@ -206,7 +206,7 @@ class LinkObservation(AbstractObservation): utilisation_category = int(utilisation_fraction * 10) + 1 # TODO: once the links support separte load per protocol, this needs amendment to reflect that. - return {"protocols": {"all": {"load": utilisation_category}}} + return {"PROTOCOLS": {"ALL": utilisation_category}} @property def space(self) -> spaces.Space: @@ -215,7 +215,7 @@ class LinkObservation(AbstractObservation): :return: Gymnasium space :rtype: spaces.Space """ - return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})}) + return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})}) @classmethod def from_config(cls, config: Dict, session: "PrimaiteSession") -> "LinkObservation": @@ -264,7 +264,6 @@ class FolderObservation(AbstractObservation): truncated_file = self.files.pop() msg = f"Too many files in folde observation. Truncating file {truncated_file}" _LOGGER.warn(msg) - raise UserWarning(msg) self.default_observation = { "health_status": 0, @@ -407,6 +406,7 @@ class NodeObservation(AbstractObservation): num_services_per_node: int = 2, num_folders_per_node: int = 2, num_files_per_folder: int = 2, + num_nics_per_node: int = 2, ) -> None: """ Configurable observation for a node in the simulation. @@ -440,18 +440,25 @@ class NodeObservation(AbstractObservation): truncated_service = self.services.pop() msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}" _LOGGER.warn(msg) - raise UserWarning(msg) # truncate service list self.folders: List[FolderObservation] = folders # add empty folder observation without `where` parameter that will always return default (blank) observations while len(self.folders) < num_folders_per_node: - self.folders.append(FolderObservation()) + self.folders.append(FolderObservation(num_files_per_folder=num_files_per_folder)) while len(self.folders) > num_folders_per_node: truncated_folder = self.folders.pop() msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}" + _LOGGER.warn(msg) self.nics: List[NicObservation] = nics + while len(self.nics) < num_nics_per_node: + self.nics.append(NicObservation()) + while len(self.nics) > num_nics_per_node: + truncated_nic = self.nics.pop() + msg = f"Too many NICs in Node observation for node. Truncating service {truncated_nic.where[-1]}" + _LOGGER.warn(msg) + self.logon_status: bool = logon_status self.default_observation: Dict = { @@ -512,6 +519,7 @@ class NodeObservation(AbstractObservation): num_services_per_node: int = 2, num_folders_per_node: int = 2, num_files_per_folder: int = 2, + num_nics_per_node: int = 2, ) -> "NodeObservation": """Create node observation from a config. Also creates child service, folder and NIC observations. @@ -562,6 +570,7 @@ class NodeObservation(AbstractObservation): num_services_per_node=num_services_per_node, num_folders_per_node=num_folders_per_node, num_files_per_folder=num_files_per_folder, + num_nics_per_node=num_nics_per_node, ) @@ -605,19 +614,17 @@ class AclObservation(AbstractObservation): self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)} "List of protocols which are part of the game, defines ordering when converting to an ID" self.default_observation: Dict = { - "RULES": { - i - + 1: { - "position": i, - "permission": 0, - "source_node_id": 0, - "source_port": 0, - "dest_node_id": 0, - "dest_port": 0, - "protocol": 0, - } - for i in range(self.num_rules) + i + + 1: { + "position": i, + "permission": 0, + "source_node_id": 0, + "source_port": 0, + "dest_node_id": 0, + "dest_port": 0, + "protocol": 0, } + for i in range(self.num_rules) } def observe(self, state: Dict) -> Dict: @@ -636,10 +643,9 @@ class AclObservation(AbstractObservation): # TODO: what if the ACL has more rules than num of max rules for obs space obs = {} - obs["RULES"] = {} for i, rule_state in acl_state.items(): if rule_state is None: - obs["RULES"][i + 1] = { + obs[i + 1] = { "position": i, "permission": 0, "source_node_id": 0, @@ -649,7 +655,7 @@ class AclObservation(AbstractObservation): "protocol": 0, } else: - obs["RULES"][i + 1] = { + obs[i + 1] = { "position": i, "permission": rule_state["action"], "source_node_id": self.node_to_id[rule_state["src_ip_address"]], @@ -669,24 +675,20 @@ class AclObservation(AbstractObservation): """ return spaces.Dict( { - "RULES": spaces.Dict( + i + + 1: spaces.Dict( { - i - + 1: spaces.Dict( - { - "position": spaces.Discrete(self.num_rules), - "permission": spaces.Discrete(3), - # adding two to lengths is to account for reserved values 0 (unused) and 1 (any) - "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), - "source_port": spaces.Discrete(len(self.port_to_id) + 2), - "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), - "dest_port": spaces.Discrete(len(self.port_to_id) + 2), - "protocol": spaces.Discrete(len(self.protocol_to_id) + 2), - } - ) - for i in range(self.num_rules) + "position": spaces.Discrete(self.num_rules), + "permission": spaces.Discrete(3), + # adding two to lengths is to account for reserved values 0 (unused) and 1 (any) + "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), + "source_port": spaces.Discrete(len(self.port_to_id) + 2), + "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), + "dest_port": spaces.Discrete(len(self.port_to_id) + 2), + "protocol": spaces.Discrete(len(self.protocol_to_id) + 2), } ) + for i in range(self.num_rules) } ) @@ -701,6 +703,7 @@ class AclObservation(AbstractObservation): :return: Observation object :rtype: AclObservation """ + max_acl_rules = config["options"]["max_acl_rules"] node_ip_to_idx = {} for ip_idx, ip_map_config in enumerate(config["ip_address_order"]): node_ref = ip_map_config["node_ref"] @@ -715,6 +718,7 @@ class AclObservation(AbstractObservation): ports=session.options.ports, protocols=session.options.protocols, where=["network", "nodes", router_uuid, "acl", "acl"], + num_rules=max_acl_rules, ) @@ -846,6 +850,7 @@ class UC2BlueObservation(AbstractObservation): num_services_per_node = config["num_services_per_node"] num_folders_per_node = config["num_folders_per_node"] num_files_per_folder = config["num_files_per_folder"] + num_nics_per_node = config["num_nics_per_node"] nodes = [ NodeObservation.from_config( config=n, @@ -853,6 +858,7 @@ class UC2BlueObservation(AbstractObservation): num_services_per_node=num_services_per_node, num_folders_per_node=num_folders_per_node, num_files_per_folder=num_files_per_folder, + num_nics_per_node=num_nics_per_node, ) for n in node_configs ] From 174de013fbd7b99afe5af662bdfa5e76d38b4aad Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 24 Oct 2023 11:43:25 +0100 Subject: [PATCH 30/35] Align blue actions with common env --- example_config.yaml | 73 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 72 insertions(+), 1 deletion(-) diff --git a/example_config.yaml b/example_config.yaml index bcc819ae..1e1150d8 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -420,13 +420,84 @@ game_config: 38: action: "NETWORK_NIC_DISABLE" options: - node_id: 6 + node_id: 1 nic_id: 1 39: action: "NETWORK_NIC_ENABLE" + options: + node_id: 1 + nic_id: 1 + 40: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 2 + nic_id: 1 + 41: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 2 + nic_id: 1 + 42: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 3 + nic_id: 1 + 43: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 3 + nic_id: 1 + 44: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 4 + nic_id: 1 + 45: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 4 + nic_id: 1 + 46: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 5 + nic_id: 1 + 47: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 5 + nic_id: 1 + 48: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 5 + nic_id: 2 + 49: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 5 + nic_id: 2 + 50: + action: "NETWORK_NIC_DISABLE" options: node_id: 6 nic_id: 1 + 51: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 6 + nic_id: 1 + 52: + action: "NETWORK_NIC_DISABLE" + options: + node_id: 7 + nic_id: 1 + 53: + action: "NETWORK_NIC_ENABLE" + options: + node_id: 7 + nic_id: 1 + options: nodes: From d49d5b292d4eac6c39108a743d1d4607f2f8bcdd Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 24 Oct 2023 15:51:29 +0100 Subject: [PATCH 31/35] Add a bit of documentation --- docs/source/config(v3).rst | 15 ++++ docs/source/high_level_project_structure.rst | 87 ++++++++++++++++++++ example_config.yaml | 8 +- 3 files changed, 106 insertions(+), 4 deletions(-) create mode 100644 docs/source/config(v3).rst create mode 100644 docs/source/high_level_project_structure.rst diff --git a/docs/source/config(v3).rst b/docs/source/config(v3).rst new file mode 100644 index 00000000..eb40015f --- /dev/null +++ b/docs/source/config(v3).rst @@ -0,0 +1,15 @@ +Primaite v3 config +****************** + + + +The YAML file allows configuring a cybersecurity scenario involving a computer network and multiple agents. There are three main sections: training_config, game, and simulation. + +The simulation section describes the simulated network environment with which the agetns interact. + +The game section describes the agents and their capabilities. Each agent has a unique type and is associated with a team (GREEN, RED, or BLUE). Each agent has a configurable observation space, action space, and reward function. + +The training_config section describes the training parameters for the learning agents. This includes the number of episodes, the number of steps per episode, and the number of steps before the agents start learning. The training_config section also describes the learning algorithm used by the agents. The learning algorithm is specified by the name of the algorithm and the hyperparameters for the algorithm. The hyperparameters are specific to each algorithm and are described in the documentation for each algorithm. + +.. only:: comment + This needs a bit of refactoring so I haven't written extensive documentation about the config yet. diff --git a/docs/source/high_level_project_structure.rst b/docs/source/high_level_project_structure.rst new file mode 100644 index 00000000..05716022 --- /dev/null +++ b/docs/source/high_level_project_structure.rst @@ -0,0 +1,87 @@ +Primaite Codebase Documentation +=============================== + +High-level structure +-------------------- +The Primaite codebase consists of two main modules: the agent-training infrastructure and the simulation logic. These modules have been decoupled to allow for flexibility and modularity. The 'game' module acts as an interface between agents and the simulation. + +Simulation +---------- +The simulation module purely simulates a computer network. It has no concept of agents acting, but it can interact with agents by providing a 'state' dictionary (using the SimComponent describe_state() method) and by accepting requests (a list of strings). + +Game layer +---------- + +The game layer is responsible for managing agents and getting them to interface with the simulator correctly. It consists of several components: + +Observations +^^^^^^^^^^^^^^^^^^ + +The ObservationManager is responsible for generating observations from the simulator state dictionary. The data is formatted so it's compatible with Gymnasium.spaces. The ObservationManager is used by the AgentManager to generate observations for each agent. + +Actions +^^^^^^^ + +The ActionManager is responsible for converting actions selected by agents (which comply with Gymnasium.spaces API) into simulation-friendly requests. The ActionManager is used by the AgentManager to take actions for each agent. + +Rewards +^^^^^^^ + +The RewardManager is responsible for calculating rewards based on the state (similar to observations). The RewardManager is used by the AgentManager to calculate rewards for each agent. + +Agents +^^^^^^ + +The AgentManager is responsible for managing agents and their interactions with the simulator. It uses the ObservationManager to generate observations for each agent, the ActionManager to take actions for each agent, and the RewardManager to calculate rewards for each agent. + +PrimaiteSession +^^^^^^^^^^^^^^^ + +PrimaiteSession is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. PrimaiteSession uses the AgentManager to manage agents and their interactions with the simulator. + +Code snippets +------------- +Here's an example of how to create a PrimaiteSession object: + +.. code-block:: python + + from primaite import PrimaiteSession + + session = PrimaiteSession() + +To start the simulation, use the start() method: + +.. code-block:: python + + session.start() + +To stop the simulation, use the stop() method: + +.. code-block:: python + + session.stop() + +To get the current state of the simulation, use the describe_state() method. This is also used as input for generating observations and rewards: + +.. code-block:: python + + state = session.sim.describe_state() + +To get the current observation of an agent, use the get_observation() method: + +.. code-block:: python + + observation = session.get_observation(agent_id) + +To get the current reward of an agent, use the get_reward() method: + +.. code-block:: python + + reward = session.get_reward(agent_id) + +To take an action for an agent, use the take_action() method: + +.. code-block:: python + + action = agent.select_action(observation) + session.take_action(agent_id, action) diff --git a/example_config.yaml b/example_config.yaml index 1e1150d8..ee42cf4f 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -2,10 +2,10 @@ training_config: rl_framework: SB3 rl_algorithm: PPO seed: 333 - n_learn_episodes: 1 - n_learn_steps: 8 - n_eval_episodes: 0 - n_eval_steps: 8 + n_learn_episodes: 20 + n_learn_steps: 128 + n_eval_episodes: 20 + n_eval_steps: 128 game_config: From cb60b6f7851f4e41b94466be90b80be2dca7698b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 24 Oct 2023 17:02:29 +0100 Subject: [PATCH 32/35] Update text in docs --- docs/index.rst | 2 + docs/source/config(v3).rst | 4 +- docs/source/game_layer.rst | 48 +++++++++++ docs/source/high_level_project_structure.rst | 87 -------------------- 4 files changed, 51 insertions(+), 90 deletions(-) create mode 100644 docs/source/game_layer.rst delete mode 100644 docs/source/high_level_project_structure.rst diff --git a/docs/index.rst b/docs/index.rst index 19f95e95..22e880fc 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -98,7 +98,9 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! source/getting_started source/about source/config + source/config(v3) source/simulation + source/game_layer source/primaite_session source/custom_agent PrimAITE API diff --git a/docs/source/config(v3).rst b/docs/source/config(v3).rst index eb40015f..0ce8b547 100644 --- a/docs/source/config(v3).rst +++ b/docs/source/config(v3).rst @@ -1,9 +1,7 @@ Primaite v3 config ****************** - - -The YAML file allows configuring a cybersecurity scenario involving a computer network and multiple agents. There are three main sections: training_config, game, and simulation. +PrimAITE uses a single configuration file to define a cybersecurity scenario. This includes the computer network and multiple agents. There are three main sections: training_config, game, and simulation. The simulation section describes the simulated network environment with which the agetns interact. diff --git a/docs/source/game_layer.rst b/docs/source/game_layer.rst new file mode 100644 index 00000000..9e254ac6 --- /dev/null +++ b/docs/source/game_layer.rst @@ -0,0 +1,48 @@ +PrimAITE Game layer +******************* + +The Primaite codebase consists of two main modules: + +* ``simulator``: The simulation logic including the network topology, the network state, and behaviour of various hardware and software classes. +* ``game``: The agent-training infrastructure which helps reinforcement learning agents interface with the simulation. This includes the observation, action, and rewards, for RL agents, but also scripted deterministic agents. The game layer orchestrates all the interactions between modules, including ARCD GATE. + +These two components have been decoupled to allow the agent training code in ARCD GATE to be reused with other simulators. The simulator and game layer communicate using the PrimAITE State API and the PrimAITE Request API. The game layer communicates with ARCD gate using the `Farama Gymnasium Spaces API `_. + +.. + TODO: write up these APIs and link them here. + + +Game layer +---------- + +The game layer is responsible for managing agents and getting them to interface with the simulator correctly. It consists of several components: + +PrimaiteSession +^^^^^^^^^^^^^^^ + +PrimaiteSession is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. PrimaiteSession keeps track of multiple agents of different types. + +Agents +^^^^^^ + +All agents inherit from the AbstractAgent class, which mandates that they have an ObservationManager, ActionManager, and RewardManager. The agent behaviour depends on the type of agent, but there are two main types: +* RL agents action during each step is decided by an RL algorithm which lives inside of ARCD GATE. The agent within PrimAITE just acts to format and forward actions decided by an RL policy. +* Deterministic agents perform all of their decision making within the PrimAITE game layer. They typically have a scripted policy which always performs the same action or a rule-based policy which performs actions based on the current state of the simulation. They can have a stochastic element, and their seed will be settable. + +.. + TODO: add seed to stochastic scripted agents + +Observations +^^^^^^^^^^^^^^^^^^ + +An agent's observations are managed by the ObservationManager class. It generates observations based on the current simulation state dictionary. It also provides the observation space during initial setup. The data is formatted so it's compatible with Gymnasium.spaces. Observation spaces are composed of one or more components which are defined by the AbstractObservation base class. + +Actions +^^^^^^^ + +An agent's actions are managed by the ActionManager. It converts actions selected by agents (which are typically integers chosen from a ``gymnasium.spaces.Discrete`` space) into simulation-friendly requests. It also provides the action space during initial setup. Action spaces are composed of one or more components which are defined by the AbstractAction base class. + +Rewards +^^^^^^^ + +An agent's reward function is managed by the RewardManager. It calculates rewards based on the simulation state (in a way similar to observations). Rewards can be defined as a weighted sum of small reward components. For example, an agents reward can be based on the uptime of a database service plus the loss rate of packets between clients and a web server. The reward components are defined by the AbstractReward base class. diff --git a/docs/source/high_level_project_structure.rst b/docs/source/high_level_project_structure.rst deleted file mode 100644 index 05716022..00000000 --- a/docs/source/high_level_project_structure.rst +++ /dev/null @@ -1,87 +0,0 @@ -Primaite Codebase Documentation -=============================== - -High-level structure --------------------- -The Primaite codebase consists of two main modules: the agent-training infrastructure and the simulation logic. These modules have been decoupled to allow for flexibility and modularity. The 'game' module acts as an interface between agents and the simulation. - -Simulation ----------- -The simulation module purely simulates a computer network. It has no concept of agents acting, but it can interact with agents by providing a 'state' dictionary (using the SimComponent describe_state() method) and by accepting requests (a list of strings). - -Game layer ----------- - -The game layer is responsible for managing agents and getting them to interface with the simulator correctly. It consists of several components: - -Observations -^^^^^^^^^^^^^^^^^^ - -The ObservationManager is responsible for generating observations from the simulator state dictionary. The data is formatted so it's compatible with Gymnasium.spaces. The ObservationManager is used by the AgentManager to generate observations for each agent. - -Actions -^^^^^^^ - -The ActionManager is responsible for converting actions selected by agents (which comply with Gymnasium.spaces API) into simulation-friendly requests. The ActionManager is used by the AgentManager to take actions for each agent. - -Rewards -^^^^^^^ - -The RewardManager is responsible for calculating rewards based on the state (similar to observations). The RewardManager is used by the AgentManager to calculate rewards for each agent. - -Agents -^^^^^^ - -The AgentManager is responsible for managing agents and their interactions with the simulator. It uses the ObservationManager to generate observations for each agent, the ActionManager to take actions for each agent, and the RewardManager to calculate rewards for each agent. - -PrimaiteSession -^^^^^^^^^^^^^^^ - -PrimaiteSession is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. PrimaiteSession uses the AgentManager to manage agents and their interactions with the simulator. - -Code snippets -------------- -Here's an example of how to create a PrimaiteSession object: - -.. code-block:: python - - from primaite import PrimaiteSession - - session = PrimaiteSession() - -To start the simulation, use the start() method: - -.. code-block:: python - - session.start() - -To stop the simulation, use the stop() method: - -.. code-block:: python - - session.stop() - -To get the current state of the simulation, use the describe_state() method. This is also used as input for generating observations and rewards: - -.. code-block:: python - - state = session.sim.describe_state() - -To get the current observation of an agent, use the get_observation() method: - -.. code-block:: python - - observation = session.get_observation(agent_id) - -To get the current reward of an agent, use the get_reward() method: - -.. code-block:: python - - reward = session.get_reward(agent_id) - -To take an action for an agent, use the take_action() method: - -.. code-block:: python - - action = agent.select_action(observation) - session.take_action(agent_id, action) From c9b06e1bfbf0fd3673575a1b1c8988a6d312812f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 24 Oct 2023 17:06:58 +0100 Subject: [PATCH 33/35] Delete sandbox file --- sandbox.ipynb | 869 -------------------------------------------------- 1 file changed, 869 deletions(-) delete mode 100644 sandbox.ipynb diff --git a/sandbox.ipynb b/sandbox.ipynb deleted file mode 100644 index 73c3e682..00000000 --- a/sandbox.ipynb +++ /dev/null @@ -1,869 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-10 21:00:40,310: Added node f0fb4743-43d5-4741-a0e5-78340a458a11 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,313: Added node 86f4ee9d-386d-4436-96e7-d00dd8bae746 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,316: Added node d2b38f6a-10fe-4d28-86a5-a1d2010c4bea to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,319: Added node 29d55bd4-d59e-4f73-95ad-b430c8716b15 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,323: Added node 5d266f99-03a6-47b3-ab0b-28b8a6813a11 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,333: Added node e4f397c7-eed2-4e69-afaf-ce44664ff1d2 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,338: Added node 19dbb2d4-648b-4387-aa15-99757b513acb to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,344: Added node f508ea65-e660-4b91-8628-5c586075137b to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,350: Added node ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,356: Added node 0ce9efc6-39ae-4d3d-82ad-43be5ac57586 to Network 301a746b-6521-4ff9-876f-a1ce39678451\n", - "2023-10-10 21:00:40,360: NIC ba:25:b9:f4:b0:74/192.168.1.1 connected to Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36\n", - "2023-10-10 21:00:40,361: SwitchPort 3a:25:73:6c:4c:36 connected to Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36\n", - "2023-10-10 21:00:40,363: Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36 up\n", - "2023-10-10 21:00:40,364: Link ba:25:b9:f4:b0:74/192.168.1.1<-->3a:25:73:6c:4c:36 up\n", - "2023-10-10 21:00:40,366: Added link eda2befd-9f68-440c-9b20-06838d9e6553 to connect ba:25:b9:f4:b0:74/192.168.1.1 and 3a:25:73:6c:4c:36\n", - "2023-10-10 21:00:40,367: NIC c0:ce:5a:6d:71:73/192.168.1.1 connected to Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3\n", - "2023-10-10 21:00:40,368: SwitchPort b2:a4:b3:fc:da:c3 connected to Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3\n", - "2023-10-10 21:00:40,370: Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3 up\n", - "2023-10-10 21:00:40,370: Link c0:ce:5a:6d:71:73/192.168.1.1<-->b2:a4:b3:fc:da:c3 up\n", - "2023-10-10 21:00:40,372: Added link 0e2a1df6-1c3c-4517-b41f-04a5cefe307c to connect c0:ce:5a:6d:71:73/192.168.1.1 and b2:a4:b3:fc:da:c3\n", - "2023-10-10 21:00:40,373: SwitchPort 72:f3:93:a5:a5:59 connected to Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10\n", - "2023-10-10 21:00:40,377: Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10 up\n", - "2023-10-10 21:00:40,379: NIC e1:9b:c0:59:1d:46/192.168.1.10 connected to Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10\n", - "2023-10-10 21:00:40,380: Link 72:f3:93:a5:a5:59<-->e1:9b:c0:59:1d:46/192.168.1.10 up\n", - "2023-10-10 21:00:40,381: Added link 48c9173d-847b-441f-9f15-f7838cf09083 to connect 72:f3:93:a5:a5:59 and e1:9b:c0:59:1d:46/192.168.1.10\n", - "2023-10-10 21:00:40,382: SwitchPort 67:36:e8:51:35:f2 connected to Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12\n", - "2023-10-10 21:00:40,384: Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12 up\n", - "2023-10-10 21:00:40,385: NIC 3b:a3:b4:ec:a0:2a/192.168.1.12 connected to Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12\n", - "2023-10-10 21:00:40,386: Link 67:36:e8:51:35:f2<-->3b:a3:b4:ec:a0:2a/192.168.1.12 up\n", - "2023-10-10 21:00:40,386: Added link 7259c2a4-44f7-44d9-87de-1a692c98ac58 to connect 67:36:e8:51:35:f2 and 3b:a3:b4:ec:a0:2a/192.168.1.12\n", - "2023-10-10 21:00:40,388: SwitchPort be:56:a8:d3:f9:6c connected to Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14\n", - "2023-10-10 21:00:40,391: Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14 up\n", - "2023-10-10 21:00:40,392: NIC 99:24:3a:ad:99:5b/192.168.1.14 connected to Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14\n", - "2023-10-10 21:00:40,394: Link be:56:a8:d3:f9:6c<-->99:24:3a:ad:99:5b/192.168.1.14 up\n", - "2023-10-10 21:00:40,396: Added link ff3c3f9a-c267-421a-9e2f-bcc11a6fa082 to connect be:56:a8:d3:f9:6c and 99:24:3a:ad:99:5b/192.168.1.14\n", - "2023-10-10 21:00:40,397: SwitchPort 4f:91:78:5a:68:3f connected to Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16\n", - "2023-10-10 21:00:40,399: Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16 up\n", - "2023-10-10 21:00:40,400: NIC 2b:ff:81:44:7f:61/192.168.1.16 connected to Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16\n", - "2023-10-10 21:00:40,401: Link 4f:91:78:5a:68:3f<-->2b:ff:81:44:7f:61/192.168.1.16 up\n", - "2023-10-10 21:00:40,402: Added link b5a2e61c-f14d-45a3-a27e-2451f21229d5 to connect 4f:91:78:5a:68:3f and 2b:ff:81:44:7f:61/192.168.1.16\n", - "2023-10-10 21:00:40,403: SwitchPort 14:8e:63:89:28:f6 connected to Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110\n", - "2023-10-10 21:00:40,404: Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110 up\n", - "2023-10-10 21:00:40,405: NIC 72:ca:83:f9:d0:66/192.168.1.110 connected to Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110\n", - "2023-10-10 21:00:40,406: Link 14:8e:63:89:28:f6<-->72:ca:83:f9:d0:66/192.168.1.110 up\n", - "2023-10-10 21:00:40,407: Added link d0906ccd-5f98-415a-b08e-b8ebfdd3c5a9 to connect 14:8e:63:89:28:f6 and 72:ca:83:f9:d0:66/192.168.1.110\n", - "2023-10-10 21:00:40,409: SwitchPort e5:20:55:91:b5:b9 connected to Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21\n", - "2023-10-10 21:00:40,412: Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21 up\n", - "2023-10-10 21:00:40,413: NIC 27:2f:d1:75:04:f9/192.168.10.21 connected to Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21\n", - "2023-10-10 21:00:40,414: Link e5:20:55:91:b5:b9<-->27:2f:d1:75:04:f9/192.168.10.21 up\n", - "2023-10-10 21:00:40,415: Added link 51f14a80-8fd6-4de9-86a3-9c726c036167 to connect e5:20:55:91:b5:b9 and 27:2f:d1:75:04:f9/192.168.10.21\n", - "2023-10-10 21:00:40,416: SwitchPort 4e:58:c4:9c:0c:6d connected to Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22\n", - "2023-10-10 21:00:40,418: Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22 up\n", - "2023-10-10 21:00:40,420: NIC 2e:f5:71:bb:73:ec/192.168.10.22 connected to Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22\n", - "2023-10-10 21:00:40,420: Link 4e:58:c4:9c:0c:6d<-->2e:f5:71:bb:73:ec/192.168.10.22 up\n", - "2023-10-10 21:00:40,421: Added link 7463a72f-af42-435f-ad98-83bcfe5728a3 to connect 4e:58:c4:9c:0c:6d and 2e:f5:71:bb:73:ec/192.168.10.22\n", - "2023-10-10 21:00:40,423: SwitchPort 66:98:bd:ca:08:b0 connected to Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110\n", - "2023-10-10 21:00:40,425: Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110 up\n", - "2023-10-10 21:00:40,427: NIC cf:f7:fc:d1:f4:ae/192.168.10.110 connected to Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110\n", - "2023-10-10 21:00:40,428: Link 66:98:bd:ca:08:b0<-->cf:f7:fc:d1:f4:ae/192.168.10.110 up\n", - "2023-10-10 21:00:40,430: Added link ff6fd52d-f893-4c0d-99d2-28f10c128043 to connect 66:98:bd:ca:08:b0 and cf:f7:fc:d1:f4:ae/192.168.10.110\n", - "2023-10-10 21:00:40,431: Stepping primaite session. Step counter: 0\n", - "2023-10-10 21:00:40,432: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,433: Getting agent action\n", - "2023-10-10 21:00:40,435: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,436: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,437: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,439: Getting agent action\n", - "2023-10-10 21:00:40,440: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,441: Sending request to simulation: ['do_nothing']\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "/home/cade/.local/state/primaite/2.0.0/log\n", - "service type not found DatabaseBackup\n", - "service type not found WebBrowser\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-10 21:00:40,443: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,445: Getting agent action\n", - "2023-10-10 21:00:40,447: Formatting agent action NODE_FOLDER_RESTORE\n", - "2023-10-10 21:00:40,449: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,450: Initiating simulation step 0\n" - ] - } - ], - "source": [ - "%load_ext autoreload\n", - "%autoreload 2\n", - "from primaite.game.session import PrimaiteSession\n", - "\n", - "from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS\n", - "import logging\n", - "_PRIMAITE_CONFIG['log_level']=logging.DEBUG\n", - "print(PRIMAITE_PATHS.app_log_dir_path)\n", - "from primaite.game.session import PrimaiteSession\n", - "import yaml\n", - "\n", - "with open('example_config.yaml', 'r') as file:\n", - " cfg = yaml.safe_load(file)\n", - "sess = PrimaiteSession.from_config(cfg)\n", - "sess.step()" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2023-10-10 21:00:40,481: Stepping primaite session. Step counter: 1\n", - "2023-10-10 21:00:40,484: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,486: Getting agent action\n", - "2023-10-10 21:00:40,487: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,488: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,489: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,491: Getting agent action\n", - "2023-10-10 21:00:40,493: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,494: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,495: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,497: Getting agent action\n", - "2023-10-10 21:00:40,498: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:40,499: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 8]\n", - "2023-10-10 21:00:40,500: Initiating simulation step 1\n", - "2023-10-10 21:00:40,502: Stepping primaite session. Step counter: 2\n", - "2023-10-10 21:00:40,503: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,504: Getting agent action\n", - "2023-10-10 21:00:40,505: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,505: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,506: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,508: Getting agent action\n", - "2023-10-10 21:00:40,511: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,513: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,513: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,515: Getting agent action\n", - "2023-10-10 21:00:40,516: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:40,518: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,519: Initiating simulation step 2\n", - "2023-10-10 21:00:40,520: Stepping primaite session. Step counter: 3\n", - "2023-10-10 21:00:40,521: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,522: Getting agent action\n", - "2023-10-10 21:00:40,524: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,525: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,526: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,528: Getting agent action\n", - "2023-10-10 21:00:40,530: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,531: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,533: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,534: Getting agent action\n", - "2023-10-10 21:00:40,535: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:40,537: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 1]\n", - "2023-10-10 21:00:40,538: Initiating simulation step 3\n", - "2023-10-10 21:00:40,539: Stepping primaite session. Step counter: 4\n", - "2023-10-10 21:00:40,541: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,543: Getting agent action\n", - "2023-10-10 21:00:40,545: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,546: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,547: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,550: Getting agent action\n", - "2023-10-10 21:00:40,552: Formatting agent action NODE_OS_SCAN\n", - "2023-10-10 21:00:40,554: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", - "2023-10-10 21:00:40,555: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,573: Getting agent action\n", - "2023-10-10 21:00:40,575: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:40,577: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 2]\n", - "2023-10-10 21:00:40,578: Initiating simulation step 4\n", - "2023-10-10 21:00:40,580: Stepping primaite session. Step counter: 5\n", - "2023-10-10 21:00:40,581: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,583: Getting agent action\n", - "2023-10-10 21:00:40,585: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,586: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,587: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,597: Getting agent action\n", - "2023-10-10 21:00:40,598: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,599: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,601: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,603: Getting agent action\n", - "2023-10-10 21:00:40,605: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:40,606: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 0]\n", - "2023-10-10 21:00:40,613: Initiating simulation step 5\n", - "2023-10-10 21:00:40,614: Stepping primaite session. Step counter: 6\n", - "2023-10-10 21:00:40,615: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,617: Getting agent action\n", - "2023-10-10 21:00:40,618: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,619: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,620: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,621: Getting agent action\n", - "2023-10-10 21:00:40,623: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,624: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,625: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,626: Getting agent action\n", - "2023-10-10 21:00:40,628: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:40,629: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 7]\n", - "2023-10-10 21:00:40,630: Initiating simulation step 6\n", - "2023-10-10 21:00:40,631: Stepping primaite session. Step counter: 7\n", - "2023-10-10 21:00:40,631: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,633: Getting agent action\n", - "2023-10-10 21:00:40,634: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,635: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,635: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,636: Getting agent action\n", - "2023-10-10 21:00:40,639: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,640: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,642: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,644: Getting agent action\n", - "2023-10-10 21:00:40,645: Formatting agent action NODE_SERVICE_STOP\n", - "2023-10-10 21:00:40,646: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,647: Initiating simulation step 7\n", - "2023-10-10 21:00:40,648: Stepping primaite session. Step counter: 8\n", - "2023-10-10 21:00:40,649: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,651: Getting agent action\n", - "2023-10-10 21:00:40,652: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,652: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,653: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,655: Getting agent action\n", - "2023-10-10 21:00:40,656: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,657: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,659: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,661: Getting agent action\n", - "2023-10-10 21:00:40,662: Formatting agent action NODE_FILE_RESTORE\n", - "2023-10-10 21:00:40,663: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,664: Initiating simulation step 8\n", - "2023-10-10 21:00:40,665: Stepping primaite session. Step counter: 9\n", - "2023-10-10 21:00:40,667: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,668: Getting agent action\n", - "2023-10-10 21:00:40,669: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,670: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,672: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,674: Getting agent action\n", - "2023-10-10 21:00:40,676: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,677: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,678: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,679: Getting agent action\n", - "2023-10-10 21:00:40,681: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,682: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,683: Initiating simulation step 9\n", - "2023-10-10 21:00:40,684: Stepping primaite session. Step counter: 10\n", - "2023-10-10 21:00:40,685: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,687: Getting agent action\n", - "2023-10-10 21:00:40,688: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,689: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,690: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,694: Getting agent action\n", - "2023-10-10 21:00:40,696: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,697: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,698: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,700: Getting agent action\n", - "2023-10-10 21:00:40,702: Formatting agent action NODE_FILE_SCAN\n", - "2023-10-10 21:00:40,705: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,706: Initiating simulation step 10\n", - "2023-10-10 21:00:40,709: Stepping primaite session. Step counter: 11\n", - "2023-10-10 21:00:40,711: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,713: Getting agent action\n", - "2023-10-10 21:00:40,715: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,716: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,717: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,719: Getting agent action\n", - "2023-10-10 21:00:40,722: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,724: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,726: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,728: Getting agent action\n", - "2023-10-10 21:00:40,730: Formatting agent action NODE_SERVICE_START\n", - "2023-10-10 21:00:40,731: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,733: Initiating simulation step 11\n", - "2023-10-10 21:00:40,735: Stepping primaite session. Step counter: 12\n", - "2023-10-10 21:00:40,736: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,738: Getting agent action\n", - "2023-10-10 21:00:40,739: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,740: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,741: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,743: Getting agent action\n", - "2023-10-10 21:00:40,746: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,748: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,749: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,752: Getting agent action\n", - "2023-10-10 21:00:40,755: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:40,758: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:40,760: Initiating simulation step 12\n", - "2023-10-10 21:00:40,762: Stepping primaite session. Step counter: 13\n", - "2023-10-10 21:00:40,763: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,766: Getting agent action\n", - "2023-10-10 21:00:40,768: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,769: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,771: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,773: Getting agent action\n", - "2023-10-10 21:00:40,777: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,779: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,780: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,782: Getting agent action\n", - "2023-10-10 21:00:40,784: Formatting agent action NODE_SERVICE_RESTART\n", - "2023-10-10 21:00:40,785: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,786: Initiating simulation step 13\n", - "2023-10-10 21:00:40,787: Stepping primaite session. Step counter: 14\n", - "2023-10-10 21:00:40,788: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,790: Getting agent action\n", - "2023-10-10 21:00:40,792: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,794: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,795: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,797: Getting agent action\n", - "2023-10-10 21:00:40,799: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,800: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,801: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,803: Getting agent action\n", - "2023-10-10 21:00:40,805: Formatting agent action NODE_SERVICE_DISABLE\n", - "2023-10-10 21:00:40,806: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,807: Initiating simulation step 14\n", - "2023-10-10 21:00:40,808: Stepping primaite session. Step counter: 15\n", - "2023-10-10 21:00:40,809: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,813: Getting agent action\n", - "2023-10-10 21:00:40,815: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,817: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,818: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,821: Getting agent action\n", - "2023-10-10 21:00:40,822: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,824: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,828: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,830: Getting agent action\n", - "2023-10-10 21:00:40,832: Formatting agent action NETWORK_NIC_DISABLE\n", - "2023-10-10 21:00:40,833: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,835: Initiating simulation step 15\n", - "2023-10-10 21:00:40,836: Stepping primaite session. Step counter: 16\n", - "2023-10-10 21:00:40,838: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,840: Getting agent action\n", - "2023-10-10 21:00:40,848: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,852: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,856: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,858: Getting agent action\n", - "2023-10-10 21:00:40,863: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,868: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,869: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,872: Getting agent action\n", - "2023-10-10 21:00:40,878: Formatting agent action NODE_OS_SCAN\n", - "2023-10-10 21:00:40,883: Sending request to simulation: ['network', 'node', '29d55bd4-d59e-4f73-95ad-b430c8716b15', 'scan']\n", - "2023-10-10 21:00:40,885: Initiating simulation step 16\n", - "2023-10-10 21:00:40,887: Stepping primaite session. Step counter: 17\n", - "2023-10-10 21:00:40,889: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,900: Getting agent action\n", - "2023-10-10 21:00:40,904: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,909: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,911: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,913: Getting agent action\n", - "2023-10-10 21:00:40,915: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,917: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,918: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,920: Getting agent action\n", - "2023-10-10 21:00:40,921: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:40,922: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:40,924: Initiating simulation step 17\n", - "2023-10-10 21:00:40,925: Stepping primaite session. Step counter: 18\n", - "2023-10-10 21:00:40,927: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,929: Getting agent action\n", - "2023-10-10 21:00:40,931: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,933: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,934: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,936: Getting agent action\n", - "2023-10-10 21:00:40,938: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,940: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,941: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,945: Getting agent action\n", - "2023-10-10 21:00:40,947: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:40,948: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,949: Initiating simulation step 18\n", - "2023-10-10 21:00:40,951: Stepping primaite session. Step counter: 19\n", - "2023-10-10 21:00:40,952: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,954: Getting agent action\n", - "2023-10-10 21:00:40,955: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,957: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,960: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,962: Getting agent action\n", - "2023-10-10 21:00:40,964: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:40,966: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,967: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,969: Getting agent action\n", - "2023-10-10 21:00:40,971: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,972: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,973: Initiating simulation step 19\n", - "2023-10-10 21:00:40,975: Stepping primaite session. Step counter: 20\n", - "2023-10-10 21:00:40,976: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:40,979: Getting agent action\n", - "2023-10-10 21:00:40,981: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:40,982: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,983: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:40,986: Getting agent action\n", - "2023-10-10 21:00:40,987: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:40,989: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:40,990: Sending simulation state to agent defender\n", - "2023-10-10 21:00:40,992: Getting agent action\n", - "2023-10-10 21:00:40,994: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:40,996: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:40,997: Initiating simulation step 20\n", - "2023-10-10 21:00:40,998: Stepping primaite session. Step counter: 21\n", - "2023-10-10 21:00:41,000: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,001: Getting agent action\n", - "2023-10-10 21:00:41,003: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,004: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,005: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,006: Getting agent action\n", - "2023-10-10 21:00:41,009: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,010: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,011: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,014: Getting agent action\n", - "2023-10-10 21:00:41,016: Formatting agent action NODE_FILE_REPAIR\n", - "2023-10-10 21:00:41,017: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,018: Initiating simulation step 21\n", - "2023-10-10 21:00:41,020: Stepping primaite session. Step counter: 22\n", - "2023-10-10 21:00:41,021: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,023: Getting agent action\n", - "2023-10-10 21:00:41,024: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,025: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,027: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,030: Getting agent action\n", - "2023-10-10 21:00:41,031: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,033: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,034: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,036: Getting agent action\n", - "2023-10-10 21:00:41,037: Formatting agent action NODE_FILE_REPAIR\n", - "2023-10-10 21:00:41,038: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,039: Initiating simulation step 22\n", - "2023-10-10 21:00:41,040: Stepping primaite session. Step counter: 23\n", - "2023-10-10 21:00:41,042: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,045: Getting agent action\n", - "2023-10-10 21:00:41,047: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,049: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,050: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,056: Getting agent action\n", - "2023-10-10 21:00:41,065: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,067: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,068: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,072: Getting agent action\n", - "2023-10-10 21:00:41,073: Formatting agent action NODE_FOLDER_REPAIR\n", - "2023-10-10 21:00:41,075: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,077: Initiating simulation step 23\n", - "2023-10-10 21:00:41,079: Stepping primaite session. Step counter: 24\n", - "2023-10-10 21:00:41,081: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,083: Getting agent action\n", - "2023-10-10 21:00:41,085: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,086: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,088: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,090: Getting agent action\n", - "2023-10-10 21:00:41,092: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,093: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,095: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,097: Getting agent action\n", - "2023-10-10 21:00:41,099: Formatting agent action NODE_STARTUP\n", - "2023-10-10 21:00:41,100: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'startup']\n", - "2023-10-10 21:00:41,102: Initiating simulation step 24\n", - "2023-10-10 21:00:41,103: Stepping primaite session. Step counter: 25\n", - "2023-10-10 21:00:41,104: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,107: Getting agent action\n", - "2023-10-10 21:00:41,110: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,112: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,113: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,114: Getting agent action\n", - "2023-10-10 21:00:41,116: Formatting agent action NODE_OS_SCAN\n", - "2023-10-10 21:00:41,118: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", - "2023-10-10 21:00:41,120: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,122: Getting agent action\n", - "2023-10-10 21:00:41,124: Formatting agent action NODE_FOLDER_RESTORE\n", - "2023-10-10 21:00:41,126: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,127: Initiating simulation step 25\n", - "2023-10-10 21:00:41,129: Stepping primaite session. Step counter: 26\n", - "2023-10-10 21:00:41,130: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,132: Getting agent action\n", - "2023-10-10 21:00:41,134: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,136: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,137: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,139: Getting agent action\n", - "2023-10-10 21:00:41,141: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,142: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,144: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,145: Getting agent action\n", - "2023-10-10 21:00:41,147: Formatting agent action NODE_STARTUP\n", - "2023-10-10 21:00:41,148: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'startup']\n", - "2023-10-10 21:00:41,150: Initiating simulation step 26\n", - "2023-10-10 21:00:41,151: Stepping primaite session. Step counter: 27\n", - "2023-10-10 21:00:41,153: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,155: Getting agent action\n", - "2023-10-10 21:00:41,157: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,158: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,160: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,162: Getting agent action\n", - "2023-10-10 21:00:41,164: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,166: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,166: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,169: Getting agent action\n", - "2023-10-10 21:00:41,170: Formatting agent action NODE_SERVICE_START\n", - "2023-10-10 21:00:41,172: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,173: Initiating simulation step 27\n", - "2023-10-10 21:00:41,174: Stepping primaite session. Step counter: 28\n", - "2023-10-10 21:00:41,176: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,177: Getting agent action\n", - "2023-10-10 21:00:41,179: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,181: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,182: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,183: Getting agent action\n", - "2023-10-10 21:00:41,185: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,186: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,187: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,188: Getting agent action\n", - "2023-10-10 21:00:41,190: Formatting agent action NETWORK_NIC_DISABLE\n", - "2023-10-10 21:00:41,191: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,193: Initiating simulation step 28\n", - "2023-10-10 21:00:41,194: Stepping primaite session. Step counter: 29\n", - "2023-10-10 21:00:41,196: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,198: Getting agent action\n", - "2023-10-10 21:00:41,199: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,201: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,201: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,203: Getting agent action\n", - "2023-10-10 21:00:41,204: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,206: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,207: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,208: Getting agent action\n", - "2023-10-10 21:00:41,211: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:41,213: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 3]\n", - "2023-10-10 21:00:41,215: Initiating simulation step 29\n", - "2023-10-10 21:00:41,217: Stepping primaite session. Step counter: 30\n", - "2023-10-10 21:00:41,218: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,220: Getting agent action\n", - "2023-10-10 21:00:41,221: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,222: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,223: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,225: Getting agent action\n", - "2023-10-10 21:00:41,227: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,229: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,231: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,232: Getting agent action\n", - "2023-10-10 21:00:41,233: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:41,235: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:41,236: Initiating simulation step 30\n", - "2023-10-10 21:00:41,238: Stepping primaite session. Step counter: 31\n", - "2023-10-10 21:00:41,238: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,240: Getting agent action\n", - "2023-10-10 21:00:41,242: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,244: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,245: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,248: Getting agent action\n", - "2023-10-10 21:00:41,249: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,251: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,252: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,254: Getting agent action\n", - "2023-10-10 21:00:41,255: Formatting agent action NODE_FILE_RESTORE\n", - "2023-10-10 21:00:41,256: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,257: Initiating simulation step 31\n", - "2023-10-10 21:00:41,259: Stepping primaite session. Step counter: 32\n", - "2023-10-10 21:00:41,260: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,263: Getting agent action\n", - "2023-10-10 21:00:41,264: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,266: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,267: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,268: Getting agent action\n", - "2023-10-10 21:00:41,269: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,270: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,272: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,273: Getting agent action\n", - "2023-10-10 21:00:41,275: Formatting agent action NODE_FILE_RESTORE\n", - "2023-10-10 21:00:41,277: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,278: Initiating simulation step 32\n", - "2023-10-10 21:00:41,280: Stepping primaite session. Step counter: 33\n", - "2023-10-10 21:00:41,281: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,284: Getting agent action\n", - "2023-10-10 21:00:41,286: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,288: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,289: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,291: Getting agent action\n", - "2023-10-10 21:00:41,293: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,295: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,299: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,301: Getting agent action\n", - "2023-10-10 21:00:41,307: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:41,309: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 2]\n", - "2023-10-10 21:00:41,311: Initiating simulation step 33\n", - "2023-10-10 21:00:41,312: Stepping primaite session. Step counter: 34\n", - "2023-10-10 21:00:41,322: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,325: Getting agent action\n", - "2023-10-10 21:00:41,327: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,328: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,339: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,341: Getting agent action\n", - "2023-10-10 21:00:41,343: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,345: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,346: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,348: Getting agent action\n", - "2023-10-10 21:00:41,350: Formatting agent action NODE_FOLDER_CHECKHASH\n", - "2023-10-10 21:00:41,352: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,354: Initiating simulation step 34\n", - "2023-10-10 21:00:41,355: Stepping primaite session. Step counter: 35\n", - "2023-10-10 21:00:41,357: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,359: Getting agent action\n", - "2023-10-10 21:00:41,360: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,362: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,363: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,366: Getting agent action\n", - "2023-10-10 21:00:41,368: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,369: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,370: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,371: Getting agent action\n", - "2023-10-10 21:00:41,373: Formatting agent action NODE_FILE_RESTORE\n", - "2023-10-10 21:00:41,378: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,381: Initiating simulation step 35\n", - "2023-10-10 21:00:41,382: Stepping primaite session. Step counter: 36\n", - "2023-10-10 21:00:41,384: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,386: Getting agent action\n", - "2023-10-10 21:00:41,388: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,389: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,390: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,392: Getting agent action\n", - "2023-10-10 21:00:41,394: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,395: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,397: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,399: Getting agent action\n", - "2023-10-10 21:00:41,401: Formatting agent action NODE_SERVICE_RESUME\n", - "2023-10-10 21:00:41,402: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,403: Initiating simulation step 36\n", - "2023-10-10 21:00:41,405: Stepping primaite session. Step counter: 37\n", - "2023-10-10 21:00:41,406: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,408: Getting agent action\n", - "2023-10-10 21:00:41,409: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,410: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,413: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,415: Getting agent action\n", - "2023-10-10 21:00:41,416: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,417: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,418: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,420: Getting agent action\n", - "2023-10-10 21:00:41,422: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:41,422: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,424: Initiating simulation step 37\n", - "2023-10-10 21:00:41,425: Stepping primaite session. Step counter: 38\n", - "2023-10-10 21:00:41,427: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,429: Getting agent action\n", - "2023-10-10 21:00:41,431: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,432: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,433: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,434: Getting agent action\n", - "2023-10-10 21:00:41,436: Formatting agent action NODE_OS_SCAN\n", - "2023-10-10 21:00:41,438: Sending request to simulation: ['network', 'node', 'ffe02a33-7f9c-4fc1-ad3a-791935dbd4c2', 'scan']\n", - "2023-10-10 21:00:41,440: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,442: Getting agent action\n", - "2023-10-10 21:00:41,444: Formatting agent action NODE_FILE_CHECKHASH\n", - "2023-10-10 21:00:41,449: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,451: Initiating simulation step 38\n", - "2023-10-10 21:00:41,452: Stepping primaite session. Step counter: 39\n", - "2023-10-10 21:00:41,453: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,457: Getting agent action\n", - "2023-10-10 21:00:41,459: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,461: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,462: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,464: Getting agent action\n", - "2023-10-10 21:00:41,465: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,467: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,469: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,472: Getting agent action\n", - "2023-10-10 21:00:41,475: Formatting agent action NODE_SERVICE_SCAN\n", - "2023-10-10 21:00:41,476: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,477: Initiating simulation step 39\n", - "2023-10-10 21:00:41,479: Stepping primaite session. Step counter: 40\n", - "2023-10-10 21:00:41,480: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,482: Getting agent action\n", - "2023-10-10 21:00:41,484: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,486: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,487: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,489: Getting agent action\n", - "2023-10-10 21:00:41,490: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,491: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,493: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,495: Getting agent action\n", - "2023-10-10 21:00:41,497: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:41,498: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 7]\n", - "2023-10-10 21:00:41,499: Initiating simulation step 40\n", - "2023-10-10 21:00:41,501: Stepping primaite session. Step counter: 41\n", - "2023-10-10 21:00:41,503: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,504: Getting agent action\n", - "2023-10-10 21:00:41,506: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,507: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,509: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,510: Getting agent action\n", - "2023-10-10 21:00:41,513: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,514: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,515: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,517: Getting agent action\n", - "2023-10-10 21:00:41,519: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:41,520: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 0]\n", - "2023-10-10 21:00:41,521: Initiating simulation step 41\n", - "2023-10-10 21:00:41,522: Stepping primaite session. Step counter: 42\n", - "2023-10-10 21:00:41,523: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,524: Getting agent action\n", - "2023-10-10 21:00:41,527: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,528: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,530: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,532: Getting agent action\n", - "2023-10-10 21:00:41,534: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,535: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,536: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,538: Getting agent action\n", - "2023-10-10 21:00:41,540: Formatting agent action NODE_RESET\n", - "2023-10-10 21:00:41,541: Sending request to simulation: ['network', 'node', '19dbb2d4-648b-4387-aa15-99757b513acb', 'reset']\n", - "2023-10-10 21:00:41,543: Initiating simulation step 42\n", - "2023-10-10 21:00:41,544: Stepping primaite session. Step counter: 43\n", - "2023-10-10 21:00:41,546: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,548: Getting agent action\n", - "2023-10-10 21:00:41,550: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,551: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,552: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,554: Getting agent action\n", - "2023-10-10 21:00:41,555: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,556: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,557: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,559: Getting agent action\n", - "2023-10-10 21:00:41,561: Formatting agent action NODE_FOLDER_RESTORE\n", - "2023-10-10 21:00:41,563: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,565: Initiating simulation step 43\n", - "2023-10-10 21:00:41,566: Stepping primaite session. Step counter: 44\n", - "2023-10-10 21:00:41,567: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,568: Getting agent action\n", - "2023-10-10 21:00:41,569: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,570: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,571: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,574: Getting agent action\n", - "2023-10-10 21:00:41,575: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,577: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,578: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,581: Getting agent action\n", - "2023-10-10 21:00:41,583: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:41,584: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.12'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:41,586: Initiating simulation step 44\n", - "2023-10-10 21:00:41,587: Stepping primaite session. Step counter: 45\n", - "2023-10-10 21:00:41,588: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,590: Getting agent action\n", - "2023-10-10 21:00:41,591: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,593: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,594: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,596: Getting agent action\n", - "2023-10-10 21:00:41,598: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,599: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,600: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,602: Getting agent action\n", - "2023-10-10 21:00:41,603: Formatting agent action NETWORK_ACL_REMOVERULE\n", - "2023-10-10 21:00:41,604: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'remove_rule', 1]\n", - "2023-10-10 21:00:41,605: Initiating simulation step 45\n", - "2023-10-10 21:00:41,606: Stepping primaite session. Step counter: 46\n", - "2023-10-10 21:00:41,607: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,608: Getting agent action\n", - "2023-10-10 21:00:41,610: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,612: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,613: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,616: Getting agent action\n", - "2023-10-10 21:00:41,617: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,618: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,620: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,621: Getting agent action\n", - "2023-10-10 21:00:41,623: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:41,624: Sending request to simulation: ['network', 'node', 'f0fb4743-43d5-4741-a0e5-78340a458a11', 'acl', 'add_rule', 'DENY', 'TCP', IPv4Address('192.168.1.10'), 'ALL', IPv4Address('127.0.0.1'), 'ALL', 1]\n", - "2023-10-10 21:00:41,626: Initiating simulation step 46\n", - "2023-10-10 21:00:41,627: Stepping primaite session. Step counter: 47\n", - "2023-10-10 21:00:41,628: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,630: Getting agent action\n", - "2023-10-10 21:00:41,632: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,634: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,635: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,636: Getting agent action\n", - "2023-10-10 21:00:41,638: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,639: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,640: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,642: Getting agent action\n", - "2023-10-10 21:00:41,643: Formatting agent action NETWORK_ACL_ADDRULE\n", - "2023-10-10 21:00:41,644: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,646: Initiating simulation step 47\n", - "2023-10-10 21:00:41,648: Stepping primaite session. Step counter: 48\n", - "2023-10-10 21:00:41,649: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,651: Getting agent action\n", - "2023-10-10 21:00:41,652: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,653: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,654: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,657: Getting agent action\n", - "2023-10-10 21:00:41,658: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,661: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,663: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,665: Getting agent action\n", - "2023-10-10 21:00:41,667: Formatting agent action NODE_SERVICE_STOP\n", - "2023-10-10 21:00:41,669: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,670: Initiating simulation step 48\n", - "2023-10-10 21:00:41,672: Stepping primaite session. Step counter: 49\n", - "2023-10-10 21:00:41,674: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,684: Getting agent action\n", - "2023-10-10 21:00:41,687: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,689: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,690: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,701: Getting agent action\n", - "2023-10-10 21:00:41,702: Formatting agent action NODE_FILE_CORRUPT\n", - "2023-10-10 21:00:41,705: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,715: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,718: Getting agent action\n", - "2023-10-10 21:00:41,720: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,734: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,737: Initiating simulation step 49\n", - "2023-10-10 21:00:41,738: Stepping primaite session. Step counter: 50\n", - "2023-10-10 21:00:41,740: Sending simulation state to agent client_1_green_user\n", - "2023-10-10 21:00:41,742: Getting agent action\n", - "2023-10-10 21:00:41,744: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,747: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,750: Sending simulation state to agent client_1_data_manipulation_red_bot\n", - "2023-10-10 21:00:41,753: Getting agent action\n", - "2023-10-10 21:00:41,755: Formatting agent action NODE_FILE_DELETE\n", - "2023-10-10 21:00:41,756: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,758: Sending simulation state to agent defender\n", - "2023-10-10 21:00:41,761: Getting agent action\n", - "2023-10-10 21:00:41,771: Formatting agent action DONOTHING\n", - "2023-10-10 21:00:41,774: Sending request to simulation: ['do_nothing']\n", - "2023-10-10 21:00:41,776: Initiating simulation step 50\n" - ] - } - ], - "source": [ - "for i in range(50):\n", - " sess.step()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "venv", - "language": "python", - "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.12" - }, - "orig_nbformat": 4 - }, - "nbformat": 4, - "nbformat_minor": 2 -} From 4872c939ff46e53eaf31e6e67ba045b807343c3d Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:19:24 +0100 Subject: [PATCH 34/35] Apply suggestions from code review --- docs/source/config(v2).rst | 489 +++++++++++++++++ docs/source/config(v3).rst | 13 - docs/source/config.rst | 500 +----------------- docs/source/game_layer.rst | 12 +- .../simulator/file_system/file_system.py | 1 - 5 files changed, 507 insertions(+), 508 deletions(-) create mode 100644 docs/source/config(v2).rst delete mode 100644 docs/source/config(v3).rst diff --git a/docs/source/config(v2).rst b/docs/source/config(v2).rst new file mode 100644 index 00000000..daf7f90b --- /dev/null +++ b/docs/source/config(v2).rst @@ -0,0 +1,489 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +.. _config: + +The Config Files Explained +========================== + +PrimAITE uses two configuration files for its operation: + +* **The Training Config** + + Used to define the top-level settings of the PrimAITE environment, the reward values, and the session that is to be run. + +* **The Lay Down Config** + + Used to define the low-level settings of a session, including the network laydown, green / red agent information exchange requirements (IERSs) and Access Control Rules. + +Training Config: +******************* + +The Training Config file consists of the following attributes: + +**Generic Config Values** + + +* **agent_framework** [enum] + + This identifies the agent framework to be used to instantiate the agent algorithm. Select from one of the following: + + * NONE - Where a user developed agent is to be used + * SB3 - Stable Baselines3 + * RLLIB - Ray RLlib. + +* **agent_identifier** + + This identifies the agent to use for the session. Select from one of the following: + + * A2C - Advantage Actor Critic + * PPO - Proximal Policy Optimization + * HARDCODED - A custom built deterministic agent + * RANDOM - A Stochastic random agent + + +* **random_red_agent** [bool] + + Determines if the session should be run with a random red agent + +* **action_type** [enum] + + Determines whether a NODE, ACL, or ANY (combined NODE & ACL) action space format is adopted for the session + + +* **OBSERVATION_SPACE** [dict] + + Allows for user to configure observation space by combining one or more observation components. List of available + components is in :py:mod:`primaite.environment.observations`. + + The observation space config item should have a ``components`` key which is a list of components. Each component + config must have a ``name`` key, and can optionally have an ``options`` key. The ``options`` are passed to the + component while it is being initialised. + + This example illustrates the correct format for the observation space config item + + .. code-block:: yaml + + observation_space: + components: + - name: NODE_LINK_TABLE + - name: NODE_STATUSES + - name: LINK_TRAFFIC_LEVELS + - name: ACCESS_CONTROL_LIST + options: + combine_service_traffic : False + quantisation_levels: 99 + + + Currently available components are: + + * :py:mod:`NODE_LINK_TABLE` this does not accept any additional options + * :py:mod:`NODE_STATUSES`, this does not accept any additional options + * :py:mod:`ACCESS_CONTROL_LIST`, this does not accept additional options + * :py:mod:`LINK_TRAFFIC_LEVELS`, this accepts the following options: + + * ``combine_service_traffic`` - whether to consider bandwidth use separately for each network protocol or combine them into a single bandwidth reading (boolean) + * ``quantisation_levels`` - how many discrete bandwidth usage levels to use for encoding. This can be an integer equal to or greater than 3. + + The other configurable item is ``flatten`` which is false by default. When set to true, the observation space is flattened (turned into a 1-D vector). You should use this if your RL agent does not natively support observation space types like ``gym.Spaces.Tuple``. + +* **num_train_episodes** [int] + + This defines the number of episodes that the agent will train for. + + +* **num_train_steps** [int] + + Determines the number of steps to run in each episode of the training session. + + +* **num_eval_episodes** [int] + + This defines the number of episodes that the agent will be evaluated over. + + +* **num_eval_steps** [int] + + Determines the number of steps to run in each episode of the evaluation session. + + +* **time_delay** [int] + + The time delay (in milliseconds) to take between each step when running a GENERIC agent session + + +* **session_type** [text] + + Type of session to be run (TRAINING, EVALUATION, or BOTH) + +* **load_agent** [bool] + + Determine whether to load an agent from file + +* **agent_load_file** [text] + + File path and file name of agent if you're loading one in + +* **observation_space_high_value** [int] + + The high value to use for values in the observation space. This is set to 1000000000 by default, and should not need changing in most cases + +* **implicit_acl_rule** [str] + + Determines which Explicit rule the ACL list has - two options are: DENY or ALLOW. + +* **max_number_acl_rules** [int] + + Sets a limit on how many ACL rules there can be in the ACL list throughout the training session. + +**Reward-Based Config Values** + +Rewards are calculated based on the difference between the current state and reference state (the 'should be' state) of the environment. + +* **Generic [all_ok]** [float] + + The score to give when the current situation (for a given component) is no different from that expected in the baseline (i.e. as though no blue or red agent actions had been undertaken) + +* **Node Hardware State [off_should_be_on]** [float] + + The score to give when the node should be on, but is off + +* **Node Hardware State [off_should_be_resetting]** [float] + + The score to give when the node should be resetting, but is off + +* **Node Hardware State [on_should_be_off]** [float] + + The score to give when the node should be off, but is on + +* **Node Hardware State [on_should_be_resetting]** [float] + + The score to give when the node should be resetting, but is on + +* **Node Hardware State [resetting_should_be_on]** [float] + + The score to give when the node should be on, but is resetting + +* **Node Hardware State [resetting_should_be_off]** [float] + + The score to give when the node should be off, but is resetting + +* **Node Hardware State [resetting]** [float] + + The score to give when the node is resetting + +* **Node Operating System or Service State [good_should_be_patching]** [float] + + The score to give when the state should be patching, but is good + +* **Node Operating System or Service State [good_should_be_compromised]** [float] + + The score to give when the state should be compromised, but is good + +* **Node Operating System or Service State [good_should_be_overwhelmed]** [float] + + The score to give when the state should be overwhelmed, but is good + +* **Node Operating System or Service State [patching_should_be_good]** [float] + + The score to give when the state should be good, but is patching + +* **Node Operating System or Service State [patching_should_be_compromised]** [float] + + The score to give when the state should be compromised, but is patching + +* **Node Operating System or Service State [patching_should_be_overwhelmed]** [float] + + The score to give when the state should be overwhelmed, but is patching + +* **Node Operating System or Service State [patching]** [float] + + The score to give when the state is patching + +* **Node Operating System or Service State [compromised_should_be_good]** [float] + + The score to give when the state should be good, but is compromised + +* **Node Operating System or Service State [compromised_should_be_patching]** [float] + + The score to give when the state should be patching, but is compromised + +* **Node Operating System or Service State [compromised_should_be_overwhelmed]** [float] + + The score to give when the state should be overwhelmed, but is compromised + +* **Node Operating System or Service State [compromised]** [float] + + The score to give when the state is compromised + +* **Node Operating System or Service State [overwhelmed_should_be_good]** [float] + + The score to give when the state should be good, but is overwhelmed + +* **Node Operating System or Service State [overwhelmed_should_be_patching]** [float] + + The score to give when the state should be patching, but is overwhelmed + +* **Node Operating System or Service State [overwhelmed_should_be_compromised]** [float] + + The score to give when the state should be compromised, but is overwhelmed + +* **Node Operating System or Service State [overwhelmed]** [float] + + The score to give when the state is overwhelmed + +* **Node File System State [good_should_be_repairing]** [float] + + The score to give when the state should be repairing, but is good + +* **Node File System State [good_should_be_restoring]** [float] + + The score to give when the state should be restoring, but is good + +* **Node File System State [good_should_be_corrupt]** [float] + + The score to give when the state should be corrupt, but is good + +* **Node File System State [good_should_be_destroyed]** [float] + + The score to give when the state should be destroyed, but is good + +* **Node File System State [repairing_should_be_good]** [float] + + The score to give when the state should be good, but is repairing + +* **Node File System State [repairing_should_be_restoring]** [float] + + The score to give when the state should be restoring, but is repairing + +* **Node File System State [repairing_should_be_corrupt]** [float] + + The score to give when the state should be corrupt, but is repairing + +* **Node File System State [repairing_should_be_destroyed]** [float] + + The score to give when the state should be destroyed, but is repairing + +* **Node File System State [repairing]** [float] + + The score to give when the state is repairing + +* **Node File System State [restoring_should_be_good]** [float] + + The score to give when the state should be good, but is restoring + +* **Node File System State [restoring_should_be_repairing]** [float] + + The score to give when the state should be repairing, but is restoring + +* **Node File System State [restoring_should_be_corrupt]** [float] + + The score to give when the state should be corrupt, but is restoring + +* **Node File System State [restoring_should_be_destroyed]** [float] + + The score to give when the state should be destroyed, but is restoring + +* **Node File System State [restoring]** [float] + + The score to give when the state is restoring + +* **Node File System State [corrupt_should_be_good]** [float] + + The score to give when the state should be good, but is corrupt + +* **Node File System State [corrupt_should_be_repairing]** [float] + + The score to give when the state should be repairing, but is corrupt + +* **Node File System State [corrupt_should_be_restoring]** [float] + + The score to give when the state should be restoring, but is corrupt + +* **Node File System State [corrupt_should_be_destroyed]** [float] + + The score to give when the state should be destroyed, but is corrupt + +* **Node File System State [corrupt]** [float] + + The score to give when the state is corrupt + +* **Node File System State [destroyed_should_be_good]** [float] + + The score to give when the state should be good, but is destroyed + +* **Node File System State [destroyed_should_be_repairing]** [float] + + The score to give when the state should be repairing, but is destroyed + +* **Node File System State [destroyed_should_be_restoring]** [float] + + The score to give when the state should be restoring, but is destroyed + +* **Node File System State [destroyed_should_be_corrupt]** [float] + + The score to give when the state should be corrupt, but is destroyed + +* **Node File System State [destroyed]** [float] + + The score to give when the state is destroyed + +* **Node File System State [scanning]** [float] + + The score to give when the state is scanning + +* **IER Status [red_ier_running]** [float] + + The score to give when a red agent IER is permitted to run + +* **IER Status [green_ier_blocked]** [float] + + The score to give when a green agent IER is prevented from running + +**Patching / Reset Durations** + +* **os_patching_duration** [int] + + The number of steps to take when patching an Operating System + +* **node_reset_duration** [int] + + The number of steps to take when resetting a node's hardware state + +* **service_patching_duration** [int] + + The number of steps to take when patching a service + +* **file_system_repairing_limit** [int]: + + The number of steps to take when repairing the file system + +* **file_system_restoring_limit** [int] + + The number of steps to take when restoring the file system + +* **file_system_scanning_limit** [int] + + The number of steps to take when scanning the file system + +* **deterministic** [bool] + + Set to true if the agent evaluation should be deterministic. Default is ``False`` + +* **seed** [int] + + Seed used in the randomisation in agent training. Default is ``None`` + +The Lay Down Config +******************* + +The lay down config file consists of the following attributes: + + +* **itemType: STEPS** [int] + +* **item_type: PORTS** [int] + + Provides a list of ports modelled in this session + +* **item_type: SERVICES** [freetext] + + Provides a list of services modelled in this session + +* **item_type: NODE** + + Defines a node included in the system laydown being simulated. It should consist of the following attributes: + + * **id** [int]: Unique ID for this YAML item + * **name** [freetext]: Human-readable name of the component + * **node_class** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services + * **node_type** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER + * **priority** [enum]: Provides a priority for each node. Can be one of P1, P2, P3, P4 or P5 (which P1 being the highest) + * **hardware_state** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING + * **ip_address** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx + * **software_state** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED + * **file_system_state** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING + * **services**: For each service associated with the node: + + * **name** [freetext]: Free-text name of the service, but must match one of the services defined for the system in the services list + * **port** [int]: Integer value of the port related to this service, but must match one of the ports defined for the system in the ports list + * **state** [enum]: The initial state of the service. Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED + +* **item_type: LINK** + + Defines a link included in the system laydown being simulated. It should consist of the following attributes: + + * **id** [int]: Unique ID for this YAML item + * **name** [freetext]: Human-readable name of the component + * **bandwidth** [int]: The bandwidth (in bits/s) of the link + * **source** [int]: The ID of the source node + * **destination** [int]: The ID of the destination node + +* **item_type: GREEN_IER** + + Defines a green agent Information Exchange Requirement (IER). It should consist of: + + * **id** [int]: Unique ID for this YAML item + * **start_step** [int]: The start step (in the episode) for this IER to begin + * **end_step** [int]: The end step (in the episode) for this IER to finish + * **load** [int]: The load (in bits/s) for this IER to apply to links + * **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list + * **port** [int]: The port that the protocol is running on. This must match a value in the ports list + * **source** [int]: The ID of the source node + * **destination** [int]: The ID of the destination node + * **mission_criticality** [enum]: The mission criticality of this IER (with 5 being highest, 1 lowest) + +* **item_type: RED_IER** + + Defines a red agent Information Exchange Requirement (IER). It should consist of: + + * **id** [int]: Unique ID for this YAML item + * **start_step** [int]: The start step (in the episode) for this IER to begin + * **end_step** [int]: The end step (in the episode) for this IER to finish + * **load** [int]: The load (in bits/s) for this IER to apply to links + * **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list + * **port** [int]: The port that the protocol is running on. This must match a value in the ports list + * **source** [int]: The ID of the source node + * **destination** [int]: The ID of the destination node + * **mission_criticality** [enum]: Not currently used. Default to 0 + +* **item_type: GREEN_POL** + + Defines a green agent pattern-of-life instruction. It should consist of: + + * **id** [int]: Unique ID for this YAML item + * **start_step** [int]: The start step (in the episode) for this PoL to begin + * **end_step** [int]: Not currently used. Default to same as start step + * **nodeId** [int]: The ID of the node to apply the PoL to + * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE + * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list + * **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) + +* **item_type: RED_POL** + + Defines a red agent pattern-of-life instruction. It should consist of: + + * **id** [int]: Unique ID for this YAML item + * **start_step** [int]: The start step (in the episode) for this PoL to begin + * **end_step** [int]: Not currently used. Default to same as start step + * **targetNodeId** [int]: The ID of the node to apply the PoL to + * **initiator** [enum]: What initiates the PoL. Can be DIRECT, IER or SERVICE + * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE + * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list + * **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state) + * **sourceNodeId** [int] The ID of the source node containing the service to check (used for SERVICE initiator) + * **sourceNodeService** [freetext]: The service on the source node to check (used for SERVICE initiator). Must match a value in the services list for this node + * **sourceNodeServiceState** [enum]: The state of the source node service to check (used for SERVICE initiator). Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED + +* **item_type: ACL_RULE** + + Defines an initial Access Control List (ACL) rule. It should consist of: + + * **id** [int]: Unique ID for this YAML item + * **permission** [enum]: Defines either an allow or deny rule. Value must be either DENY or ALLOW + * **source** [IP address]: Defines the source IP address for the rule in xxx.xxx.xxx.xxx format + * **destination** [IP address]: Defines the destination IP address for the rule in xxx.xxx.xxx.xxx format + * **protocol** [freetext]: Defines the protocol for the rule. Must match a value in the services list + * **port** [int]: Defines the port for the rule. Must match a value in the ports list + * **position** [int]: Defines where to place the ACL rule in the list. Lower index or (higher up in the list) means they are checked first. Index starts at 0 (Python indexes). diff --git a/docs/source/config(v3).rst b/docs/source/config(v3).rst deleted file mode 100644 index 0ce8b547..00000000 --- a/docs/source/config(v3).rst +++ /dev/null @@ -1,13 +0,0 @@ -Primaite v3 config -****************** - -PrimAITE uses a single configuration file to define a cybersecurity scenario. This includes the computer network and multiple agents. There are three main sections: training_config, game, and simulation. - -The simulation section describes the simulated network environment with which the agetns interact. - -The game section describes the agents and their capabilities. Each agent has a unique type and is associated with a team (GREEN, RED, or BLUE). Each agent has a configurable observation space, action space, and reward function. - -The training_config section describes the training parameters for the learning agents. This includes the number of episodes, the number of steps per episode, and the number of steps before the agents start learning. The training_config section also describes the learning algorithm used by the agents. The learning algorithm is specified by the name of the algorithm and the hyperparameters for the algorithm. The hyperparameters are specific to each algorithm and are described in the documentation for each algorithm. - -.. only:: comment - This needs a bit of refactoring so I haven't written extensive documentation about the config yet. diff --git a/docs/source/config.rst b/docs/source/config.rst index daf7f90b..0ce8b547 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -1,489 +1,13 @@ +Primaite v3 config +****************** + +PrimAITE uses a single configuration file to define a cybersecurity scenario. This includes the computer network and multiple agents. There are three main sections: training_config, game, and simulation. + +The simulation section describes the simulated network environment with which the agetns interact. + +The game section describes the agents and their capabilities. Each agent has a unique type and is associated with a team (GREEN, RED, or BLUE). Each agent has a configurable observation space, action space, and reward function. + +The training_config section describes the training parameters for the learning agents. This includes the number of episodes, the number of steps per episode, and the number of steps before the agents start learning. The training_config section also describes the learning algorithm used by the agents. The learning algorithm is specified by the name of the algorithm and the hyperparameters for the algorithm. The hyperparameters are specific to each algorithm and are described in the documentation for each algorithm. + .. only:: comment - - © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -.. _config: - -The Config Files Explained -========================== - -PrimAITE uses two configuration files for its operation: - -* **The Training Config** - - Used to define the top-level settings of the PrimAITE environment, the reward values, and the session that is to be run. - -* **The Lay Down Config** - - Used to define the low-level settings of a session, including the network laydown, green / red agent information exchange requirements (IERSs) and Access Control Rules. - -Training Config: -******************* - -The Training Config file consists of the following attributes: - -**Generic Config Values** - - -* **agent_framework** [enum] - - This identifies the agent framework to be used to instantiate the agent algorithm. Select from one of the following: - - * NONE - Where a user developed agent is to be used - * SB3 - Stable Baselines3 - * RLLIB - Ray RLlib. - -* **agent_identifier** - - This identifies the agent to use for the session. Select from one of the following: - - * A2C - Advantage Actor Critic - * PPO - Proximal Policy Optimization - * HARDCODED - A custom built deterministic agent - * RANDOM - A Stochastic random agent - - -* **random_red_agent** [bool] - - Determines if the session should be run with a random red agent - -* **action_type** [enum] - - Determines whether a NODE, ACL, or ANY (combined NODE & ACL) action space format is adopted for the session - - -* **OBSERVATION_SPACE** [dict] - - Allows for user to configure observation space by combining one or more observation components. List of available - components is in :py:mod:`primaite.environment.observations`. - - The observation space config item should have a ``components`` key which is a list of components. Each component - config must have a ``name`` key, and can optionally have an ``options`` key. The ``options`` are passed to the - component while it is being initialised. - - This example illustrates the correct format for the observation space config item - - .. code-block:: yaml - - observation_space: - components: - - name: NODE_LINK_TABLE - - name: NODE_STATUSES - - name: LINK_TRAFFIC_LEVELS - - name: ACCESS_CONTROL_LIST - options: - combine_service_traffic : False - quantisation_levels: 99 - - - Currently available components are: - - * :py:mod:`NODE_LINK_TABLE` this does not accept any additional options - * :py:mod:`NODE_STATUSES`, this does not accept any additional options - * :py:mod:`ACCESS_CONTROL_LIST`, this does not accept additional options - * :py:mod:`LINK_TRAFFIC_LEVELS`, this accepts the following options: - - * ``combine_service_traffic`` - whether to consider bandwidth use separately for each network protocol or combine them into a single bandwidth reading (boolean) - * ``quantisation_levels`` - how many discrete bandwidth usage levels to use for encoding. This can be an integer equal to or greater than 3. - - The other configurable item is ``flatten`` which is false by default. When set to true, the observation space is flattened (turned into a 1-D vector). You should use this if your RL agent does not natively support observation space types like ``gym.Spaces.Tuple``. - -* **num_train_episodes** [int] - - This defines the number of episodes that the agent will train for. - - -* **num_train_steps** [int] - - Determines the number of steps to run in each episode of the training session. - - -* **num_eval_episodes** [int] - - This defines the number of episodes that the agent will be evaluated over. - - -* **num_eval_steps** [int] - - Determines the number of steps to run in each episode of the evaluation session. - - -* **time_delay** [int] - - The time delay (in milliseconds) to take between each step when running a GENERIC agent session - - -* **session_type** [text] - - Type of session to be run (TRAINING, EVALUATION, or BOTH) - -* **load_agent** [bool] - - Determine whether to load an agent from file - -* **agent_load_file** [text] - - File path and file name of agent if you're loading one in - -* **observation_space_high_value** [int] - - The high value to use for values in the observation space. This is set to 1000000000 by default, and should not need changing in most cases - -* **implicit_acl_rule** [str] - - Determines which Explicit rule the ACL list has - two options are: DENY or ALLOW. - -* **max_number_acl_rules** [int] - - Sets a limit on how many ACL rules there can be in the ACL list throughout the training session. - -**Reward-Based Config Values** - -Rewards are calculated based on the difference between the current state and reference state (the 'should be' state) of the environment. - -* **Generic [all_ok]** [float] - - The score to give when the current situation (for a given component) is no different from that expected in the baseline (i.e. as though no blue or red agent actions had been undertaken) - -* **Node Hardware State [off_should_be_on]** [float] - - The score to give when the node should be on, but is off - -* **Node Hardware State [off_should_be_resetting]** [float] - - The score to give when the node should be resetting, but is off - -* **Node Hardware State [on_should_be_off]** [float] - - The score to give when the node should be off, but is on - -* **Node Hardware State [on_should_be_resetting]** [float] - - The score to give when the node should be resetting, but is on - -* **Node Hardware State [resetting_should_be_on]** [float] - - The score to give when the node should be on, but is resetting - -* **Node Hardware State [resetting_should_be_off]** [float] - - The score to give when the node should be off, but is resetting - -* **Node Hardware State [resetting]** [float] - - The score to give when the node is resetting - -* **Node Operating System or Service State [good_should_be_patching]** [float] - - The score to give when the state should be patching, but is good - -* **Node Operating System or Service State [good_should_be_compromised]** [float] - - The score to give when the state should be compromised, but is good - -* **Node Operating System or Service State [good_should_be_overwhelmed]** [float] - - The score to give when the state should be overwhelmed, but is good - -* **Node Operating System or Service State [patching_should_be_good]** [float] - - The score to give when the state should be good, but is patching - -* **Node Operating System or Service State [patching_should_be_compromised]** [float] - - The score to give when the state should be compromised, but is patching - -* **Node Operating System or Service State [patching_should_be_overwhelmed]** [float] - - The score to give when the state should be overwhelmed, but is patching - -* **Node Operating System or Service State [patching]** [float] - - The score to give when the state is patching - -* **Node Operating System or Service State [compromised_should_be_good]** [float] - - The score to give when the state should be good, but is compromised - -* **Node Operating System or Service State [compromised_should_be_patching]** [float] - - The score to give when the state should be patching, but is compromised - -* **Node Operating System or Service State [compromised_should_be_overwhelmed]** [float] - - The score to give when the state should be overwhelmed, but is compromised - -* **Node Operating System or Service State [compromised]** [float] - - The score to give when the state is compromised - -* **Node Operating System or Service State [overwhelmed_should_be_good]** [float] - - The score to give when the state should be good, but is overwhelmed - -* **Node Operating System or Service State [overwhelmed_should_be_patching]** [float] - - The score to give when the state should be patching, but is overwhelmed - -* **Node Operating System or Service State [overwhelmed_should_be_compromised]** [float] - - The score to give when the state should be compromised, but is overwhelmed - -* **Node Operating System or Service State [overwhelmed]** [float] - - The score to give when the state is overwhelmed - -* **Node File System State [good_should_be_repairing]** [float] - - The score to give when the state should be repairing, but is good - -* **Node File System State [good_should_be_restoring]** [float] - - The score to give when the state should be restoring, but is good - -* **Node File System State [good_should_be_corrupt]** [float] - - The score to give when the state should be corrupt, but is good - -* **Node File System State [good_should_be_destroyed]** [float] - - The score to give when the state should be destroyed, but is good - -* **Node File System State [repairing_should_be_good]** [float] - - The score to give when the state should be good, but is repairing - -* **Node File System State [repairing_should_be_restoring]** [float] - - The score to give when the state should be restoring, but is repairing - -* **Node File System State [repairing_should_be_corrupt]** [float] - - The score to give when the state should be corrupt, but is repairing - -* **Node File System State [repairing_should_be_destroyed]** [float] - - The score to give when the state should be destroyed, but is repairing - -* **Node File System State [repairing]** [float] - - The score to give when the state is repairing - -* **Node File System State [restoring_should_be_good]** [float] - - The score to give when the state should be good, but is restoring - -* **Node File System State [restoring_should_be_repairing]** [float] - - The score to give when the state should be repairing, but is restoring - -* **Node File System State [restoring_should_be_corrupt]** [float] - - The score to give when the state should be corrupt, but is restoring - -* **Node File System State [restoring_should_be_destroyed]** [float] - - The score to give when the state should be destroyed, but is restoring - -* **Node File System State [restoring]** [float] - - The score to give when the state is restoring - -* **Node File System State [corrupt_should_be_good]** [float] - - The score to give when the state should be good, but is corrupt - -* **Node File System State [corrupt_should_be_repairing]** [float] - - The score to give when the state should be repairing, but is corrupt - -* **Node File System State [corrupt_should_be_restoring]** [float] - - The score to give when the state should be restoring, but is corrupt - -* **Node File System State [corrupt_should_be_destroyed]** [float] - - The score to give when the state should be destroyed, but is corrupt - -* **Node File System State [corrupt]** [float] - - The score to give when the state is corrupt - -* **Node File System State [destroyed_should_be_good]** [float] - - The score to give when the state should be good, but is destroyed - -* **Node File System State [destroyed_should_be_repairing]** [float] - - The score to give when the state should be repairing, but is destroyed - -* **Node File System State [destroyed_should_be_restoring]** [float] - - The score to give when the state should be restoring, but is destroyed - -* **Node File System State [destroyed_should_be_corrupt]** [float] - - The score to give when the state should be corrupt, but is destroyed - -* **Node File System State [destroyed]** [float] - - The score to give when the state is destroyed - -* **Node File System State [scanning]** [float] - - The score to give when the state is scanning - -* **IER Status [red_ier_running]** [float] - - The score to give when a red agent IER is permitted to run - -* **IER Status [green_ier_blocked]** [float] - - The score to give when a green agent IER is prevented from running - -**Patching / Reset Durations** - -* **os_patching_duration** [int] - - The number of steps to take when patching an Operating System - -* **node_reset_duration** [int] - - The number of steps to take when resetting a node's hardware state - -* **service_patching_duration** [int] - - The number of steps to take when patching a service - -* **file_system_repairing_limit** [int]: - - The number of steps to take when repairing the file system - -* **file_system_restoring_limit** [int] - - The number of steps to take when restoring the file system - -* **file_system_scanning_limit** [int] - - The number of steps to take when scanning the file system - -* **deterministic** [bool] - - Set to true if the agent evaluation should be deterministic. Default is ``False`` - -* **seed** [int] - - Seed used in the randomisation in agent training. Default is ``None`` - -The Lay Down Config -******************* - -The lay down config file consists of the following attributes: - - -* **itemType: STEPS** [int] - -* **item_type: PORTS** [int] - - Provides a list of ports modelled in this session - -* **item_type: SERVICES** [freetext] - - Provides a list of services modelled in this session - -* **item_type: NODE** - - Defines a node included in the system laydown being simulated. It should consist of the following attributes: - - * **id** [int]: Unique ID for this YAML item - * **name** [freetext]: Human-readable name of the component - * **node_class** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services - * **node_type** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER - * **priority** [enum]: Provides a priority for each node. Can be one of P1, P2, P3, P4 or P5 (which P1 being the highest) - * **hardware_state** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING - * **ip_address** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx - * **software_state** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED - * **file_system_state** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING - * **services**: For each service associated with the node: - - * **name** [freetext]: Free-text name of the service, but must match one of the services defined for the system in the services list - * **port** [int]: Integer value of the port related to this service, but must match one of the ports defined for the system in the ports list - * **state** [enum]: The initial state of the service. Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED - -* **item_type: LINK** - - Defines a link included in the system laydown being simulated. It should consist of the following attributes: - - * **id** [int]: Unique ID for this YAML item - * **name** [freetext]: Human-readable name of the component - * **bandwidth** [int]: The bandwidth (in bits/s) of the link - * **source** [int]: The ID of the source node - * **destination** [int]: The ID of the destination node - -* **item_type: GREEN_IER** - - Defines a green agent Information Exchange Requirement (IER). It should consist of: - - * **id** [int]: Unique ID for this YAML item - * **start_step** [int]: The start step (in the episode) for this IER to begin - * **end_step** [int]: The end step (in the episode) for this IER to finish - * **load** [int]: The load (in bits/s) for this IER to apply to links - * **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list - * **port** [int]: The port that the protocol is running on. This must match a value in the ports list - * **source** [int]: The ID of the source node - * **destination** [int]: The ID of the destination node - * **mission_criticality** [enum]: The mission criticality of this IER (with 5 being highest, 1 lowest) - -* **item_type: RED_IER** - - Defines a red agent Information Exchange Requirement (IER). It should consist of: - - * **id** [int]: Unique ID for this YAML item - * **start_step** [int]: The start step (in the episode) for this IER to begin - * **end_step** [int]: The end step (in the episode) for this IER to finish - * **load** [int]: The load (in bits/s) for this IER to apply to links - * **protocol** [freetext]: The protocol to apply to the links. This must match a value in the services list - * **port** [int]: The port that the protocol is running on. This must match a value in the ports list - * **source** [int]: The ID of the source node - * **destination** [int]: The ID of the destination node - * **mission_criticality** [enum]: Not currently used. Default to 0 - -* **item_type: GREEN_POL** - - Defines a green agent pattern-of-life instruction. It should consist of: - - * **id** [int]: Unique ID for this YAML item - * **start_step** [int]: The start step (in the episode) for this PoL to begin - * **end_step** [int]: Not currently used. Default to same as start step - * **nodeId** [int]: The ID of the node to apply the PoL to - * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE - * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list - * **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) - -* **item_type: RED_POL** - - Defines a red agent pattern-of-life instruction. It should consist of: - - * **id** [int]: Unique ID for this YAML item - * **start_step** [int]: The start step (in the episode) for this PoL to begin - * **end_step** [int]: Not currently used. Default to same as start step - * **targetNodeId** [int]: The ID of the node to apply the PoL to - * **initiator** [enum]: What initiates the PoL. Can be DIRECT, IER or SERVICE - * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE - * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list - * **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state) - * **sourceNodeId** [int] The ID of the source node containing the service to check (used for SERVICE initiator) - * **sourceNodeService** [freetext]: The service on the source node to check (used for SERVICE initiator). Must match a value in the services list for this node - * **sourceNodeServiceState** [enum]: The state of the source node service to check (used for SERVICE initiator). Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED - -* **item_type: ACL_RULE** - - Defines an initial Access Control List (ACL) rule. It should consist of: - - * **id** [int]: Unique ID for this YAML item - * **permission** [enum]: Defines either an allow or deny rule. Value must be either DENY or ALLOW - * **source** [IP address]: Defines the source IP address for the rule in xxx.xxx.xxx.xxx format - * **destination** [IP address]: Defines the destination IP address for the rule in xxx.xxx.xxx.xxx format - * **protocol** [freetext]: Defines the protocol for the rule. Must match a value in the services list - * **port** [int]: Defines the port for the rule. Must match a value in the ports list - * **position** [int]: Defines where to place the ACL rule in the list. Lower index or (higher up in the list) means they are checked first. Index starts at 0 (Python indexes). + This needs a bit of refactoring so I haven't written extensive documentation about the config yet. diff --git a/docs/source/game_layer.rst b/docs/source/game_layer.rst index 9e254ac6..27905c85 100644 --- a/docs/source/game_layer.rst +++ b/docs/source/game_layer.rst @@ -17,15 +17,15 @@ Game layer The game layer is responsible for managing agents and getting them to interface with the simulator correctly. It consists of several components: -PrimaiteSession +PrimAITE Session ^^^^^^^^^^^^^^^ -PrimaiteSession is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. PrimaiteSession keeps track of multiple agents of different types. +``PrimaiteSession`` is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. ``PrimaiteSession`` keeps track of multiple agents of different types. Agents ^^^^^^ -All agents inherit from the AbstractAgent class, which mandates that they have an ObservationManager, ActionManager, and RewardManager. The agent behaviour depends on the type of agent, but there are two main types: +All agents inherit from the :py:class:`primaite.game.agent.interface.AbstractAgent` class, which mandates that they have an ObservationManager, ActionManager, and RewardManager. The agent behaviour depends on the type of agent, but there are two main types: * RL agents action during each step is decided by an RL algorithm which lives inside of ARCD GATE. The agent within PrimAITE just acts to format and forward actions decided by an RL policy. * Deterministic agents perform all of their decision making within the PrimAITE game layer. They typically have a scripted policy which always performs the same action or a rule-based policy which performs actions based on the current state of the simulation. They can have a stochastic element, and their seed will be settable. @@ -35,14 +35,14 @@ All agents inherit from the AbstractAgent class, which mandates that they have a Observations ^^^^^^^^^^^^^^^^^^ -An agent's observations are managed by the ObservationManager class. It generates observations based on the current simulation state dictionary. It also provides the observation space during initial setup. The data is formatted so it's compatible with Gymnasium.spaces. Observation spaces are composed of one or more components which are defined by the AbstractObservation base class. +An agent's observations are managed by the ``ObservationManager`` class. It generates observations based on the current simulation state dictionary. It also provides the observation space during initial setup. The data is formatted so it's compatible with ``Gymnasium.spaces``. Observation spaces are composed of one or more components which are defined by the ``AbstractObservation`` base class. Actions ^^^^^^^ -An agent's actions are managed by the ActionManager. It converts actions selected by agents (which are typically integers chosen from a ``gymnasium.spaces.Discrete`` space) into simulation-friendly requests. It also provides the action space during initial setup. Action spaces are composed of one or more components which are defined by the AbstractAction base class. +An agent's actions are managed by the ``ActionManager``. It converts actions selected by agents (which are typically integers chosen from a ``gymnasium.spaces.Discrete`` space) into simulation-friendly requests. It also provides the action space during initial setup. Action spaces are composed of one or more components which are defined by the ``AbstractAction`` base class. Rewards ^^^^^^^ -An agent's reward function is managed by the RewardManager. It calculates rewards based on the simulation state (in a way similar to observations). Rewards can be defined as a weighted sum of small reward components. For example, an agents reward can be based on the uptime of a database service plus the loss rate of packets between clients and a web server. The reward components are defined by the AbstractReward base class. +An agent's reward function is managed by the ``RewardManager``. It calculates rewards based on the simulation state (in a way similar to observations). Rewards can be defined as a weighted sum of small reward components. For example, an agents reward can be based on the uptime of a database service plus the loss rate of packets between clients and a web server. The reward components are defined by the AbstractReward base class. diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 5d0dbedf..30cda446 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -74,7 +74,6 @@ class FileSystemItemABC(SimComponent): name: str "The name of the FileSystemItemABC." - health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD "Actual status of the current FileSystemItem" From 02901a7c99e2bc5bfe7e9bf281c51a8aee746d84 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 19:07:45 +0100 Subject: [PATCH 35/35] Apply suggestions from code review. --- docs/index.rst | 2 +- docs/source/config(v2).rst | 2 ++ src/primaite/game/agent/actions.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 22e880fc..ed66797d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -98,7 +98,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! source/getting_started source/about source/config - source/config(v3) + source/config(v2) source/simulation source/game_layer source/primaite_session diff --git a/docs/source/config(v2).rst b/docs/source/config(v2).rst index daf7f90b..35233cf5 100644 --- a/docs/source/config(v2).rst +++ b/docs/source/config(v2).rst @@ -7,6 +7,8 @@ The Config Files Explained ========================== +Note: This file describes the config files used in legacy PrimAITE v2.0. This file will be removed soon. + PrimAITE uses two configuration files for its operation: * **The Training Config** diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 0a380487..b06013cd 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -697,7 +697,7 @@ class ActionManager: def get_action(self, action: int) -> Tuple[str, Dict]: """Produce action in CAOS format.""" """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" - """The caos format is basically a action identifier, followed by parameters stored in a dictionary""" + """The CAOS format is basically a action identifier, followed by parameters stored in a dictionary""" act_identifier, act_options = self.action_map[action] return act_identifier, act_options