From f54f278fca6e71d48c5bf375a03933b8db200891 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 21 Sep 2023 10:13:01 +0100 Subject: [PATCH 01/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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 f861b1897a73a9085a85e0e5a1c450ce0fad9716 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 13:36:55 +0100 Subject: [PATCH 34/53] Remove deprecated code from v2 --- src/primaite/acl/__init__.py | 2 - src/primaite/acl/access_control_list.py | 198 --- src/primaite/acl/acl_rule.py | 87 - src/primaite/agents/__init__.py | 2 - src/primaite/agents/agent_abc.py | 319 ---- src/primaite/agents/hardcoded_abc.py | 118 -- src/primaite/agents/hardcoded_acl.py | 515 ------ src/primaite/agents/hardcoded_node.py | 125 -- src/primaite/agents/rllib.py | 287 ---- src/primaite/agents/sb3.py | 206 --- src/primaite/agents/simple.py | 59 - src/primaite/agents/utils.py | 450 ------ src/primaite/common/__init__.py | 2 - src/primaite/common/custom_typing.py | 8 - src/primaite/common/enums.py | 208 --- src/primaite/common/protocol.py | 47 - src/primaite/common/service.py | 28 - src/primaite/config/__init__.py | 2 - .../lay_down_config_1_DDOS_basic.yaml | 166 -- .../lay_down_config_2_DDOS_basic.yaml | 366 ----- .../lay_down_config_3_DOS_very_basic.yaml | 164 -- .../lay_down_config_5_data_manipulation.yaml | 546 ------- .../training/training_config_main.yaml | 168 -- src/primaite/config/lay_down_config.py | 141 -- src/primaite/config/training_config.py | 438 ----- src/primaite/data_viz/__init__.py | 15 - src/primaite/data_viz/session_plots.py | 73 - src/primaite/environment/__init__.py | 2 - src/primaite/environment/observations.py | 735 --------- src/primaite/environment/primaite_env.py | 1408 ----------------- src/primaite/environment/reward.py | 386 ----- src/primaite/links/__init__.py | 2 - src/primaite/links/link.py | 114 -- src/primaite/nodes/__init__.py | 2 - src/primaite/nodes/active_node.py | 208 --- src/primaite/nodes/node.py | 79 - .../nodes/node_state_instruction_green.py | 94 -- .../nodes/node_state_instruction_red.py | 143 -- src/primaite/nodes/passive_node.py | 42 - src/primaite/nodes/service_node.py | 190 --- src/primaite/pol/__init__.py | 2 - src/primaite/pol/green_pol.py | 264 ---- src/primaite/pol/ier.py | 147 -- src/primaite/pol/red_agent_pol.py | 353 ----- src/primaite/primaite_session.py | 228 --- .../setup/old_installation_clean_up.py | 14 - src/primaite/transactions/__init__.py | 2 - src/primaite/transactions/transaction.py | 102 -- 48 files changed, 9257 deletions(-) delete mode 100644 src/primaite/acl/__init__.py delete mode 100644 src/primaite/acl/access_control_list.py delete mode 100644 src/primaite/acl/acl_rule.py delete mode 100644 src/primaite/agents/__init__.py delete mode 100644 src/primaite/agents/agent_abc.py delete mode 100644 src/primaite/agents/hardcoded_abc.py delete mode 100644 src/primaite/agents/hardcoded_acl.py delete mode 100644 src/primaite/agents/hardcoded_node.py delete mode 100644 src/primaite/agents/rllib.py delete mode 100644 src/primaite/agents/sb3.py delete mode 100644 src/primaite/agents/simple.py delete mode 100644 src/primaite/agents/utils.py delete mode 100644 src/primaite/common/__init__.py delete mode 100644 src/primaite/common/custom_typing.py delete mode 100644 src/primaite/common/enums.py delete mode 100644 src/primaite/common/protocol.py delete mode 100644 src/primaite/common/service.py delete mode 100644 src/primaite/config/__init__.py delete mode 100644 src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml delete mode 100644 src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml delete mode 100644 src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml delete mode 100644 src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml delete mode 100644 src/primaite/config/_package_data/training/training_config_main.yaml delete mode 100644 src/primaite/config/lay_down_config.py delete mode 100644 src/primaite/config/training_config.py delete mode 100644 src/primaite/data_viz/__init__.py delete mode 100644 src/primaite/data_viz/session_plots.py delete mode 100644 src/primaite/environment/__init__.py delete mode 100644 src/primaite/environment/observations.py delete mode 100644 src/primaite/environment/primaite_env.py delete mode 100644 src/primaite/environment/reward.py delete mode 100644 src/primaite/links/__init__.py delete mode 100644 src/primaite/links/link.py delete mode 100644 src/primaite/nodes/__init__.py delete mode 100644 src/primaite/nodes/active_node.py delete mode 100644 src/primaite/nodes/node.py delete mode 100644 src/primaite/nodes/node_state_instruction_green.py delete mode 100644 src/primaite/nodes/node_state_instruction_red.py delete mode 100644 src/primaite/nodes/passive_node.py delete mode 100644 src/primaite/nodes/service_node.py delete mode 100644 src/primaite/pol/__init__.py delete mode 100644 src/primaite/pol/green_pol.py delete mode 100644 src/primaite/pol/ier.py delete mode 100644 src/primaite/pol/red_agent_pol.py delete mode 100644 src/primaite/primaite_session.py delete mode 100644 src/primaite/setup/old_installation_clean_up.py delete mode 100644 src/primaite/transactions/__init__.py delete mode 100644 src/primaite/transactions/transaction.py diff --git a/src/primaite/acl/__init__.py b/src/primaite/acl/__init__.py deleted file mode 100644 index 6dc02583..00000000 --- a/src/primaite/acl/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Access Control List. Models firewall functionality.""" diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py deleted file mode 100644 index 88943f8f..00000000 --- a/src/primaite/acl/access_control_list.py +++ /dev/null @@ -1,198 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""A class that implements the access control list implementation for the network.""" -import logging -from typing import Dict, Final, List, Union - -from primaite.acl.acl_rule import ACLRule -from primaite.common.enums import RulePermissionType - -_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) - - -class AccessControlList: - """Access Control List class.""" - - def __init__(self, implicit_permission: RulePermissionType, max_acl_rules: int) -> None: - """Init.""" - # Implicit ALLOW or DENY firewall spec - self.acl_implicit_permission = implicit_permission - # Implicit rule in ACL list - if self.acl_implicit_permission == RulePermissionType.DENY: - self.acl_implicit_rule = ACLRule(RulePermissionType.DENY, "ANY", "ANY", "ANY", "ANY") - elif self.acl_implicit_permission == RulePermissionType.ALLOW: - self.acl_implicit_rule = ACLRule(RulePermissionType.ALLOW, "ANY", "ANY", "ANY", "ANY") - else: - raise ValueError(f"implicit permission must be ALLOW or DENY, got {self.acl_implicit_permission}") - - # Maximum number of ACL Rules in ACL - self.max_acl_rules: int = max_acl_rules - # A list of ACL Rules - self._acl: List[Union[ACLRule, None]] = [None] * (self.max_acl_rules - 1) - - @property - def acl(self) -> List[Union[ACLRule, None]]: - """Public access method for private _acl.""" - return self._acl + [self.acl_implicit_rule] - - def check_address_match(self, _rule: ACLRule, _source_ip_address: str, _dest_ip_address: str) -> bool: - """Checks for IP address matches. - - :param _rule: The rule object to check - :type _rule: ACLRule - :param _source_ip_address: Source IP address to compare - :type _source_ip_address: str - :param _dest_ip_address: Destination IP address to compare - :type _dest_ip_address: str - :return: True if there is a match, otherwise False. - :rtype: bool - """ - if ( - (_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == _dest_ip_address) - or (_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == _dest_ip_address) - or (_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == "ANY") - or (_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == "ANY") - ): - return True - else: - return False - - def is_blocked(self, _source_ip_address: str, _dest_ip_address: str, _protocol: str, _port: str) -> bool: - """ - Checks for rules that block a protocol / port. - - Args: - _source_ip_address: the source IP address to check - _dest_ip_address: the destination IP address to check - _protocol: the protocol to check - _port: the port to check - - Returns: - Indicates block if all conditions are satisfied. - """ - for rule in self.acl: - if isinstance(rule, ACLRule): - if self.check_address_match(rule, _source_ip_address, _dest_ip_address): - if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY") and ( - str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" - ): - # There's a matching rule. Get the permission - if rule.get_permission() == RulePermissionType.DENY: - return True - elif rule.get_permission() == RulePermissionType.ALLOW: - return False - - # If there has been no rule to allow the IER through, it will return a blocked signal by default - return True - - def add_rule( - self, - _permission: RulePermissionType, - _source_ip: str, - _dest_ip: str, - _protocol: str, - _port: str, - _position: str, - ) -> None: - """ - Adds a new rule. - - Args: - _permission: the permission value (e.g. "ALLOW" or "DENY") - _source_ip: the source IP address - _dest_ip: the destination IP address - _protocol: the protocol - _port: the port - _position: position to insert ACL rule into ACL list (starting from index 1 and NOT 0) - """ - try: - position_index = int(_position) - except TypeError: - _LOGGER.info(f"Position {_position} could not be converted to integer.") - return - - new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - # Checks position is in correct range - if self.max_acl_rules - 1 > position_index > -1: - try: - _LOGGER.info(f"Position {position_index} is valid.") - # Check to see Agent will not overwrite current ACL in ACL list - if self._acl[position_index] is None: - _LOGGER.info(f"Inserting rule {new_rule} at position {position_index}") - # Adds rule - self._acl[position_index] = new_rule - else: - # Cannot overwrite it - _LOGGER.info(f"Error: inserting rule at non-empty position {position_index}") - return - except Exception: - _LOGGER.info(f"New Rule could NOT be added to list at position {position_index}.") - else: - _LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule") - - def remove_rule( - self, _permission: RulePermissionType, _source_ip: str, _dest_ip: str, _protocol: str, _port: str - ) -> None: - """ - Removes a rule. - - Args: - _permission: the permission value (e.g. "ALLOW" or "DENY") - _source_ip: the source IP address - _dest_ip: the destination IP address - _protocol: the protocol - _port: the port - """ - rule_to_delete = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - delete_rule_hash = hash(rule_to_delete) - - for index in range(0, len(self._acl)): - if isinstance(self._acl[index], ACLRule) and hash(self._acl[index]) == delete_rule_hash: - self._acl[index] = None - - def remove_all_rules(self) -> None: - """Removes all rules.""" - for i in range(len(self._acl)): - self._acl[i] = None - - def get_dictionary_hash( - self, _permission: RulePermissionType, _source_ip: str, _dest_ip: str, _protocol: str, _port: str - ) -> int: - """ - Produces a hash value for a rule. - - Args: - _permission: the permission value (e.g. "ALLOW" or "DENY") - _source_ip: the source IP address - _dest_ip: the destination IP address - _protocol: the protocol - _port: the port - - Returns: - Hash value based on rule parameters. - """ - rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - hash_value = hash(rule) - return hash_value - - def get_relevant_rules( - self, _source_ip_address: str, _dest_ip_address: str, _protocol: str, _port: str - ) -> Dict[int, ACLRule]: - """Get all ACL rules that relate to the given arguments. - - :param _source_ip_address: the source IP address to check - :param _dest_ip_address: the destination IP address to check - :param _protocol: the protocol to check - :param _port: the port to check - :return: Dictionary of all ACL rules that relate to the given arguments - :rtype: Dict[int, ACLRule] - """ - relevant_rules = {} - for rule in self.acl: - if self.check_address_match(rule, _source_ip_address, _dest_ip_address): - if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY" or _protocol == "ANY") and ( - str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" or str(_port) == "ANY" - ): - # There's a matching rule. - relevant_rules[self._acl.index(rule)] = rule - - return relevant_rules diff --git a/src/primaite/acl/acl_rule.py b/src/primaite/acl/acl_rule.py deleted file mode 100644 index 9c8deacd..00000000 --- a/src/primaite/acl/acl_rule.py +++ /dev/null @@ -1,87 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""A class that implements an access control list rule.""" -from primaite.common.enums import RulePermissionType - - -class ACLRule: - """Access Control List Rule class.""" - - def __init__( - self, _permission: RulePermissionType, _source_ip: str, _dest_ip: str, _protocol: str, _port: str - ) -> None: - """ - Initialise an ACL Rule. - - :param _permission: The permission (ALLOW or DENY) - :param _source_ip: The source IP address - :param _dest_ip: The destination IP address - :param _protocol: The rule protocol - :param _port: The rule port - """ - self.permission: RulePermissionType = _permission - self.source_ip: str = _source_ip - self.dest_ip: str = _dest_ip - self.protocol: str = _protocol - self.port: str = _port - - def __hash__(self) -> int: - """ - Override the hash function. - - Returns: - Returns hash of core parameters. - """ - return hash( - ( - self.permission, - self.source_ip, - self.dest_ip, - self.protocol, - self.port, - ) - ) - - def get_permission(self) -> str: - """ - Gets the permission attribute. - - Returns: - Returns permission attribute - """ - return self.permission - - def get_source_ip(self) -> str: - """ - Gets the source IP address attribute. - - Returns: - Returns source IP address attribute - """ - return self.source_ip - - def get_dest_ip(self) -> str: - """ - Gets the desintation IP address attribute. - - Returns: - Returns destination IP address attribute - """ - return self.dest_ip - - def get_protocol(self) -> str: - """ - Gets the protocol attribute. - - Returns: - Returns protocol attribute - """ - return self.protocol - - def get_port(self) -> str: - """ - Gets the port attribute. - - Returns: - Returns port attribute - """ - return self.port diff --git a/src/primaite/agents/__init__.py b/src/primaite/agents/__init__.py deleted file mode 100644 index c742daf3..00000000 --- a/src/primaite/agents/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Common interface between RL agents from different libraries and PrimAITE.""" diff --git a/src/primaite/agents/agent_abc.py b/src/primaite/agents/agent_abc.py deleted file mode 100644 index 359790ad..00000000 --- a/src/primaite/agents/agent_abc.py +++ /dev/null @@ -1,319 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from __future__ import annotations - -import json -from abc import ABC, abstractmethod -from datetime import datetime -from logging import Logger -from pathlib import Path -from typing import Any, Dict, Optional, Union -from uuid import uuid4 - -import primaite -from primaite import getLogger, PRIMAITE_PATHS -from primaite.config import lay_down_config, training_config -from primaite.config.training_config import TrainingConfig -from primaite.data_viz.session_plots import plot_av_reward_per_episode -from primaite.environment.primaite_env import Primaite -from primaite.utils.session_metadata_parser import parse_session_metadata - -_LOGGER: Logger = getLogger(__name__) - - -def get_session_path(session_timestamp: datetime) -> Path: - """ - Get the directory path the session will output to. - - This is set in the format of: - ~/primaite/2.0.0/sessions//_. - - :param session_timestamp: This is the datetime that the session started. - :return: The session directory path. - """ - date_dir = session_timestamp.strftime("%Y-%m-%d") - session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S") - session_path = PRIMAITE_PATHS.user_sessions_path / date_dir / session_path - session_path.mkdir(exist_ok=True, parents=True) - - return session_path - - -class AgentSessionABC(ABC): - """ - An ABC that manages training and/or evaluation of agents in PrimAITE. - - This class cannot be directly instantiated and must be inherited from with all implemented abstract methods - implemented. - """ - - @abstractmethod - def __init__( - self, - training_config_path: Optional[Union[str, Path]] = None, - lay_down_config_path: Optional[Union[str, Path]] = None, - session_path: Optional[Union[str, Path]] = None, - legacy_training_config: bool = False, - legacy_lay_down_config: bool = False, - ) -> None: - """ - Initialise an agent session from config files, or load a previous session. - - If training configuration and laydown configuration are provided with a session path, - the session path will be used. - - :param training_config_path: YAML file containing configurable items defined in - `primaite.config.training_config.TrainingConfig` - :type training_config_path: Union[path, str] - :param lay_down_config_path: YAML file containing configurable items for generating network laydown. - :type lay_down_config_path: Union[path, str] - :param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :param session_path: directory path of the session to load - """ - # initialise variables - self._env: Primaite - self._agent = None - self._can_learn: bool = False - self._can_evaluate: bool = False - self.is_eval = False - self.legacy_training_config = legacy_training_config - self.legacy_lay_down_config = legacy_lay_down_config - - self.session_timestamp: datetime = datetime.now() - - # convert session to path - if session_path is not None: - if not isinstance(session_path, Path): - session_path = Path(session_path) - - # if a session path is provided, load it - if not session_path.exists(): - raise Exception(f"Session could not be loaded. Path does not exist: {session_path}") - - # load session - self.load(session_path) - else: - # set training config path - if not isinstance(training_config_path, Path): - training_config_path = Path(training_config_path) - self._training_config_path: Union[Path, str] = training_config_path - self._training_config: TrainingConfig = training_config.load( - self._training_config_path, legacy_file=legacy_training_config - ) - - if not isinstance(lay_down_config_path, Path): - lay_down_config_path = Path(lay_down_config_path) - self._lay_down_config_path: Union[Path, str] = lay_down_config_path - self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config) - self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level - - # set random UUID for session - self._uuid = str(uuid4()) - "The session timestamp" - self.session_path = get_session_path(self.session_timestamp) - "The Session path" - - @property - def timestamp_str(self) -> str: - """The session timestamp as a string.""" - return self.session_timestamp.strftime("%Y-%m-%d_%H-%M-%S") - - @property - def learning_path(self) -> Path: - """The learning outputs path.""" - path = self.session_path / "learning" - path.mkdir(exist_ok=True, parents=True) - return path - - @property - def evaluation_path(self) -> Path: - """The evaluation outputs path.""" - path = self.session_path / "evaluation" - path.mkdir(exist_ok=True, parents=True) - return path - - @property - def checkpoints_path(self) -> Path: - """The Session checkpoints path.""" - path = self.learning_path / "checkpoints" - path.mkdir(exist_ok=True, parents=True) - return path - - @property - def uuid(self) -> str: - """The Agent Session UUID.""" - return self._uuid - - def _write_session_metadata_file(self) -> None: - """ - Write the ``session_metadata.json`` file. - - Creates a ``session_metadata.json`` in the ``session_path`` directory - and adds the following key/value pairs: - - - uuid: The UUID assigned to the session upon instantiation. - - start_datetime: The date & time the session started in iso format. - - end_datetime: NULL. - - total_episodes: NULL. - - total_time_steps: NULL. - - env: - - training_config: - - All training config items - - lay_down_config: - - All lay down config items - - """ - metadata_dict = { - "uuid": self.uuid, - "start_datetime": self.session_timestamp.isoformat(), - "end_datetime": None, - "learning": {"total_episodes": None, "total_time_steps": None}, - "evaluation": {"total_episodes": None, "total_time_steps": None}, - "env": { - "training_config": self._training_config.to_dict(json_serializable=True), - "lay_down_config": self._lay_down_config, - }, - } - filepath = self.session_path / "session_metadata.json" - _LOGGER.debug(f"Writing Session Metadata file: {filepath}") - with open(filepath, "w") as file: - json.dump(metadata_dict, file) - _LOGGER.debug("Finished writing session metadata file") - - def _update_session_metadata_file(self) -> None: - """ - Update the ``session_metadata.json`` file. - - Updates the `session_metadata.json`` in the ``session_path`` directory - with the following key/value pairs: - - - end_datetime: The date & time the session ended in iso format. - - total_episodes: The total number of training episodes completed. - - total_time_steps: The total number of training time steps completed. - """ - with open(self.session_path / "session_metadata.json", "r") as file: - metadata_dict = json.load(file) - - metadata_dict["end_datetime"] = datetime.now().isoformat() - if not self.is_eval: - metadata_dict["learning"]["total_episodes"] = self._env.actual_episode_count # noqa - metadata_dict["learning"]["total_time_steps"] = self._env.total_step_count # noqa - else: - metadata_dict["evaluation"]["total_episodes"] = self._env.actual_episode_count # noqa - metadata_dict["evaluation"]["total_time_steps"] = self._env.total_step_count # noqa - - filepath = self.session_path / "session_metadata.json" - _LOGGER.debug(f"Updating Session Metadata file: {filepath}") - with open(filepath, "w") as file: - json.dump(metadata_dict, file) - _LOGGER.debug("Finished updating session metadata file") - - @abstractmethod - def _setup(self) -> None: - _LOGGER.info( - "Welcome to the Primary-level AI Training Environment " f"(PrimAITE) (version: {primaite.__version__})" - ) - _LOGGER.info(f"The output directory for this session is: {self.session_path}") - self._write_session_metadata_file() - self._can_learn = True - self._can_evaluate = False - - @abstractmethod - def _save_checkpoint(self) -> None: - pass - - @abstractmethod - def learn( - self, - **kwargs: Any, - ) -> None: - """ - Train the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - if self._can_learn: - _LOGGER.info("Finished learning") - _LOGGER.debug("Writing transactions") - self._update_session_metadata_file() - self._can_evaluate = True - self.is_eval = False - - @abstractmethod - def evaluate( - self, - **kwargs: Any, - ) -> None: - """ - Evaluate the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - if self._can_evaluate: - self._update_session_metadata_file() - self.is_eval = True - self._plot_av_reward_per_episode(learning_session=False) - _LOGGER.info("Finished evaluation") - - @abstractmethod - def _get_latest_checkpoint(self) -> None: - pass - - def load(self, path: Union[str, Path]) -> None: - """Load an agent from file.""" - md_dict, training_config_path, laydown_config_path = parse_session_metadata(path) - - # set training config path - self._training_config_path: Union[Path, str] = training_config_path - self._training_config: TrainingConfig = training_config.load(self._training_config_path) - self._lay_down_config_path: Union[Path, str] = laydown_config_path - self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path) - self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level - - # set random UUID for session - self._uuid = md_dict["uuid"] - - # set the session path - self.session_path = path - "The Session path" - - @property - def _saved_agent_path(self) -> Path: - file_name = f"{self._training_config.agent_framework}_" f"{self._training_config.agent_identifier}" f".zip" - return self.learning_path / file_name - - @abstractmethod - def save(self) -> None: - """Save the agent.""" - pass - - @abstractmethod - def export(self) -> None: - """Export the agent to transportable file format.""" - pass - - def close(self) -> None: - """Closes the agent.""" - self._env.episode_av_reward_writer.close() # noqa - self._env.transaction_writer.close() # noqa - - def _plot_av_reward_per_episode(self, learning_session: bool = True) -> None: - # self.close() - title = f"PrimAITE Session {self.timestamp_str} " - subtitle = str(self._training_config) - csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv" - image_file = f"average_reward_per_episode_{self.timestamp_str}.png" - if learning_session: - title += "(Learning)" - path = self.learning_path / csv_file - image_path = self.learning_path / image_file - else: - title += "(Evaluation)" - path = self.evaluation_path / csv_file - image_path = self.evaluation_path / image_file - - fig = plot_av_reward_per_episode(path, title, subtitle) - fig.write_image(image_path) - _LOGGER.debug(f"Saved average rewards per episode plot to: {path}") diff --git a/src/primaite/agents/hardcoded_abc.py b/src/primaite/agents/hardcoded_abc.py deleted file mode 100644 index e75edbc5..00000000 --- a/src/primaite/agents/hardcoded_abc.py +++ /dev/null @@ -1,118 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import time -from abc import abstractmethod -from pathlib import Path -from typing import Any, Optional, Union - -import numpy as np - -from primaite import getLogger -from primaite.agents.agent_abc import AgentSessionABC -from primaite.environment.primaite_env import Primaite - -_LOGGER = getLogger(__name__) - - -class HardCodedAgentSessionABC(AgentSessionABC): - """ - An Agent Session ABC for evaluation deterministic agents. - - This class cannot be directly instantiated and must be inherited from with all implemented abstract methods - implemented. - """ - - def __init__( - self, - training_config_path: Optional[Union[str, Path]] = "", - lay_down_config_path: Optional[Union[str, Path]] = "", - session_path: Optional[Union[str, Path]] = None, - ) -> None: - """ - Initialise a hardcoded agent session. - - :param training_config_path: YAML file containing configurable items defined in - `primaite.config.training_config.TrainingConfig` - :type training_config_path: Union[path, str] - :param lay_down_config_path: YAML file containing configurable items for generating network laydown. - :type lay_down_config_path: Union[path, str] - """ - super().__init__(training_config_path, lay_down_config_path, session_path) - self._setup() - - def _setup(self) -> None: - self._env: Primaite = Primaite( - training_config_path=self._training_config_path, - lay_down_config_path=self._lay_down_config_path, - session_path=self.session_path, - timestamp_str=self.timestamp_str, - ) - super()._setup() - self._can_learn = False - self._can_evaluate = True - - def _save_checkpoint(self) -> None: - pass - - def _get_latest_checkpoint(self) -> None: - pass - - def learn( - self, - **kwargs: Any, - ) -> None: - """ - Train the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - _LOGGER.warning("Deterministic agents cannot learn") - - @abstractmethod - def _calculate_action(self, obs: np.ndarray) -> None: - pass - - def evaluate( - self, - **kwargs: Any, - ) -> None: - """ - Evaluate the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - self._env.set_as_eval() # noqa - self.is_eval = True - - time_steps = self._training_config.num_eval_steps - episodes = self._training_config.num_eval_episodes - - obs = self._env.reset() - for episode in range(episodes): - # Reset env and collect initial observation - for step in range(time_steps): - # Calculate action - action = self._calculate_action(obs) - - # Perform the step - obs, reward, done, info = self._env.step(action) - - if done: - break - - # Introduce a delay between steps - time.sleep(self._training_config.time_delay / 1000) - obs = self._env.reset() - self._env.close() - - @classmethod - def load(cls, path: Union[str, Path] = None) -> None: - """Load an agent from file.""" - _LOGGER.warning("Deterministic agents cannot be loaded") - - def save(self) -> None: - """Save the agent.""" - _LOGGER.warning("Deterministic agents cannot be saved") - - def export(self) -> None: - """Export the agent to transportable file format.""" - _LOGGER.warning("Deterministic agents cannot be exported") diff --git a/src/primaite/agents/hardcoded_acl.py b/src/primaite/agents/hardcoded_acl.py deleted file mode 100644 index 2440da06..00000000 --- a/src/primaite/agents/hardcoded_acl.py +++ /dev/null @@ -1,515 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from typing import Dict, List, Union - -import numpy as np - -from primaite.acl.access_control_list import AccessControlList -from primaite.acl.acl_rule import ACLRule -from primaite.agents.hardcoded_abc import HardCodedAgentSessionABC -from primaite.agents.utils import ( - get_new_action, - get_node_of_ip, - transform_action_acl_enum, - transform_change_obs_readable, -) -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import HardCodedAgentView -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.service_node import ServiceNode -from primaite.pol.ier import IER - - -class HardCodedACLAgent(HardCodedAgentSessionABC): - """An Agent Session class that implements a deterministic ACL agent.""" - - def _calculate_action(self, obs: np.ndarray) -> int: - if self._training_config.hard_coded_agent_view == HardCodedAgentView.BASIC: - # Basic view action using only the current observation - return self._calculate_action_basic_view(obs) - else: - # full view action using observation space, action - # history and reward feedback - return self._calculate_action_full_view(obs) - - def get_blocked_green_iers( - self, green_iers: Dict[str, IER], acl: AccessControlList, nodes: Dict[str, NodeUnion] - ) -> Dict[str, IER]: - """Get blocked green IERs. - - :param green_iers: Green IERs to check for being - :type green_iers: Dict[str, IER] - :param acl: Firewall rules - :type acl: AccessControlList - :param nodes: Nodes in the network - :type nodes: Dict[str,NodeUnion] - :return: Same as `green_iers` input dict, but filtered to only contain the blocked ones. - :rtype: Dict[str, IER] - """ - blocked_green_iers = {} - - for green_ier_id, green_ier in green_iers.items(): - source_node_id = green_ier.get_source_node_id() - source_node_address = nodes[source_node_id].ip_address - dest_node_id = green_ier.get_dest_node_id() - dest_node_address = nodes[dest_node_id].ip_address - protocol = green_ier.get_protocol() # e.g. 'TCP' - port = green_ier.get_port() - - # Can be blocked by an ACL or by default (no allow rule exists) - if acl.is_blocked(source_node_address, dest_node_address, protocol, port): - blocked_green_iers[green_ier_id] = green_ier - - return blocked_green_iers - - def get_matching_acl_rules_for_ier( - self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion] - ) -> Dict[int, ACLRule]: - """Get list of ACL rules which are relevant to an IER. - - :param ier: Information Exchange Request to query against the ACL list - :type ier: IER - :param acl: Firewall rules - :type acl: AccessControlList - :param nodes: Nodes in the network - :type nodes: Dict[str,NodeUnion] - :return: _description_ - :rtype: _type_ - """ - source_node_id = ier.get_source_node_id() - source_node_address = nodes[source_node_id].ip_address - dest_node_id = ier.get_dest_node_id() - dest_node_address = nodes[dest_node_id].ip_address - protocol = ier.get_protocol() # e.g. 'TCP' - port = ier.get_port() - matching_rules = acl.get_relevant_rules(source_node_address, dest_node_address, protocol, port) - return matching_rules - - def get_blocking_acl_rules_for_ier( - self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion] - ) -> Dict[int, ACLRule]: - """ - Get blocking ACL rules for an IER. - - .. warning:: - Can return empty dict but IER can still be blocked by default - (No ALLOW rule, therefore blocked). - - :param ier: Information Exchange Request to query against the ACL list - :type ier: IER - :param acl: Firewall rules - :type acl: AccessControlList - :param nodes: Nodes in the network - :type nodes: Dict[str,NodeUnion] - :return: _description_ - :rtype: _type_ - """ - matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes) - - blocked_rules = {} - for rule_key, rule_value in matching_rules.items(): - if rule_value.get_permission() == "DENY": - blocked_rules[rule_key] = rule_value - - return blocked_rules - - def get_allow_acl_rules_for_ier( - self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion] - ) -> Dict[int, ACLRule]: - """Get all allowing ACL rules for an IER. - - :param ier: Information Exchange Request to query against the ACL list - :type ier: IER - :param acl: Firewall rules - :type acl: AccessControlList - :param nodes: Nodes in the network - :type nodes: Dict[str,NodeUnion] - :return: _description_ - :rtype: _type_ - """ - matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes) - - allowed_rules = {} - for rule_key, rule_value in matching_rules.items(): - if rule_value.get_permission() == "ALLOW": - allowed_rules[rule_key] = rule_value - - return allowed_rules - - def get_matching_acl_rules( - self, - source_node_id: str, - dest_node_id: str, - protocol: str, - port: str, - acl: AccessControlList, - nodes: Dict[str, Union[ServiceNode, ActiveNode]], - services_list: List[str], - ) -> Dict[int, ACLRule]: - """Filter ACL rules to only those which are relevant to the specified nodes. - - :param source_node_id: Source node - :type source_node_id: str - :param dest_node_id: Destination nodes - :type dest_node_id: str - :param protocol: Network protocol - :type protocol: str - :param port: Network port - :type port: str - :param acl: Access Control list which will be filtered - :type acl: AccessControlList - :param nodes: The environment's node directory. - :type nodes: Dict[str, Union[ServiceNode, ActiveNode]] - :param services_list: List of services registered for the environment. - :type services_list: List[str] - :return: Filtered version of 'acl' - :rtype: Dict[str, ACLRule] - """ - if source_node_id != "ANY": - source_node_address = nodes[str(source_node_id)].ip_address - else: - source_node_address = source_node_id - - if dest_node_id != "ANY": - dest_node_address = nodes[str(dest_node_id)].ip_address - else: - dest_node_address = dest_node_id - - if protocol != "ANY": - protocol = services_list[protocol - 1] # -1 as dont have to account for ANY in list of services - # TODO: This should throw an error because protocol is a string - - matching_rules = acl.get_relevant_rules(source_node_address, dest_node_address, protocol, port) - return matching_rules - - def get_allow_acl_rules( - self, - source_node_id: int, - dest_node_id: str, - protocol: int, - port: str, - acl: AccessControlList, - nodes: Dict[str, NodeUnion], - services_list: List[str], - ) -> Dict[int, ACLRule]: - """List ALLOW rules relating to specified nodes. - - :param source_node_id: Source node id - :type source_node_id: int - :param dest_node_id: Destination node - :type dest_node_id: str - :param protocol: Network protocol - :type protocol: int - :param port: Port - :type port: str - :param acl: Firewall ruleset which is applied to the network - :type acl: AccessControlList - :param nodes: The simulation's node store - :type nodes: Dict[str, NodeUnion] - :param services_list: Services list - :type services_list: List[str] - :return: Filtered ACL Rule directory which includes only those rules which affect the specified source and - desination nodes - :rtype: Dict[str, ACLRule] - """ - matching_rules = self.get_matching_acl_rules( - source_node_id, - dest_node_id, - protocol, - port, - acl, - nodes, - services_list, - ) - - allowed_rules = {} - for rule_key, rule_value in matching_rules.items(): - if rule_value.get_permission() == "ALLOW": - allowed_rules[rule_key] = rule_value - - return allowed_rules - - def get_deny_acl_rules( - self, - source_node_id: int, - dest_node_id: str, - protocol: int, - port: str, - acl: AccessControlList, - nodes: Dict[str, NodeUnion], - services_list: List[str], - ) -> Dict[int, ACLRule]: - """List DENY rules relating to specified nodes. - - :param source_node_id: Source node id - :type source_node_id: int - :param dest_node_id: Destination node - :type dest_node_id: str - :param protocol: Network protocol - :type protocol: int - :param port: Port - :type port: str - :param acl: Firewall ruleset which is applied to the network - :type acl: AccessControlList - :param nodes: The simulation's node store - :type nodes: Dict[str, NodeUnion] - :param services_list: Services list - :type services_list: List[str] - :return: Filtered ACL Rule directory which includes only those rules which affect the specified source and - desination nodes - :rtype: Dict[str, ACLRule] - """ - matching_rules = self.get_matching_acl_rules( - source_node_id, - dest_node_id, - protocol, - port, - acl, - nodes, - services_list, - ) - - allowed_rules = {} - for rule_key, rule_value in matching_rules.items(): - if rule_value.get_permission() == "DENY": - allowed_rules[rule_key] = rule_value - - return allowed_rules - - def _calculate_action_full_view(self, obs: np.ndarray) -> int: - """ - Calculate a good acl-based action for the blue agent to take. - - Knowledge of just the observation space is insufficient for a perfect solution, as we need to know: - - - Which ACL rules already exist, - otherwise: - - The agent would perminently get stuck in a loop of performing the same action over and over. - (best action is to block something, but its already blocked but doesn't know this) - - The agent would be unable to interact with existing rules (e.g. how would it know to delete a rule, - if it doesnt know what rules exist) - - The Green IERs (optional) - It often needs to know which traffic it should be allowing. For example - in the default config one of the green IERs is blocked by default, but it has no way of knowing this - based on the observation space. Additionally, potentially in the future, once a node state - has been fixed (no longer compromised), it needs a way to know it should reallow traffic. - A RL agent can learn what the green IERs are on its own - but the rule based agent cannot easily do this. - - There doesn't seem like there's much that can be done if an Operating or OS State is compromised - - If a service node becomes compromised there's a decision to make - do we block that service? - Pros: It cannot launch an attack on another node, so the node will not be able to be OVERWHELMED - Cons: Will block a green IER, decreasing the reward - We decide to block the service. - - Potentially a better solution (for the reward) would be to block the incomming traffic from compromised - nodes once a service becomes overwhelmed. However currently the ACL action space has no way of reversing - an overwhelmed state, so we don't do this. - - :param obs: current observation from the gym environment - :type obs: np.ndarray - :return: Optimal action to take in the environment (chosen from the discrete action space) - :rtype: int - """ - # obs = convert_to_old_obs(obs) - r_obs = transform_change_obs_readable(obs) - _, _, _, *s = r_obs - - if len(r_obs) == 4: # only 1 service - s = [*s] - - # 1. Check if node is compromised. If so we want to block its outwards services - # a. If it is comprimised check if there's an allow rule we should delete. - # cons: might delete a multi-rule from any source node (ANY -> x) - # b. OPTIONAL (Deny rules not needed): Check if there already exists an existing Deny Rule so not to duplicate - # c. OPTIONAL (no allow rule = blocked): Add a DENY rule - found_action = False - for service_num, service_states in enumerate(s): - for x, service_state in enumerate(service_states): - if service_state == "COMPROMISED": - action_source_id = x + 1 # +1 as 0 is any - action_destination_id = "ANY" - action_protocol = service_num + 1 # +1 as 0 is any - action_port = "ANY" - - allow_rules = self.get_allow_acl_rules( - action_source_id, - action_destination_id, - action_protocol, - action_port, - self._env.acl, - self._env.nodes, - self._env.services_list, - ) - deny_rules = self.get_deny_acl_rules( - action_source_id, - action_destination_id, - action_protocol, - action_port, - self._env.acl, - self._env.nodes, - self._env.services_list, - ) - if len(allow_rules) > 0: - # Check if there's an allow rule we should delete - rule = list(allow_rules.values())[0] - action_decision = "DELETE" - action_permission = "ALLOW" - action_source_ip = rule.get_source_ip() - action_source_id = int(get_node_of_ip(action_source_ip, self._env.nodes)) - action_destination_ip = rule.get_dest_ip() - action_destination_id = int(get_node_of_ip(action_destination_ip, self._env.nodes)) - action_protocol_name = rule.get_protocol() - action_protocol = ( - self._env.services_list.index(action_protocol_name) + 1 - ) # convert name e.g. 'TCP' to index - action_port_name = rule.get_port() - action_port = ( - self._env.ports_list.index(action_port_name) + 1 - ) # convert port name e.g. '80' to index - - found_action = True - break - elif len(deny_rules) > 0: - # TODO OPTIONAL - # If there's already a DENY RULE, that blocks EVERYTHING from the source ip we don't need - # to create another - # Check to see if the DENY rule really blocks everything (ANY) or just a specific rule - continue - else: - # TODO OPTIONAL: Add a DENY rule, optional as by default no allow rule == blocked - action_decision = "CREATE" - action_permission = "DENY" - break - if found_action: - break - - # 2. If NO Node is Comprimised, or the node has already been blocked, check the green IERs and - # add an Allow rule if the green IER is being blocked. - # a. OPTIONAL - NOT IMPLEMENTED (optional as a deny rule does not overwrite an allow rule): - # If there's a DENY rule delete it if: - # - There isn't already a deny rule - # - It doesnt allows a comprimised node to become operational. - # b. Add an ALLOW rule if: - # - There isn't already an allow rule - # - It doesnt allows a comprimised node to become operational - - if not found_action: - # Which Green IERS are blocked - blocked_green_iers = self.get_blocked_green_iers(self._env.green_iers, self._env.acl, self._env.nodes) - for ier_key, ier in blocked_green_iers.items(): - # Which ALLOW rules are allowing this IER (none) - allowing_rules = self.get_allow_acl_rules_for_ier(ier, self._env.acl, self._env.nodes) - - # If there are no blocking rules, it may be being blocked by default - # If there is already an allow rule - node_id_to_check = int(ier.get_source_node_id()) - service_name_to_check = ier.get_protocol() - service_id_to_check = self._env.services_list.index(service_name_to_check) - - # Service state of the the source node in the ier - service_state = s[service_id_to_check][node_id_to_check - 1] - - if len(allowing_rules) == 0 and service_state != "COMPROMISED": - action_decision = "CREATE" - action_permission = "ALLOW" - action_source_id = int(ier.get_source_node_id()) - action_destination_id = int(ier.get_dest_node_id()) - action_protocol_name = ier.get_protocol() - action_protocol = ( - self._env.services_list.index(action_protocol_name) + 1 - ) # convert name e.g. 'TCP' to index - action_port_name = ier.get_port() - action_port = ( - self._env.ports_list.index(action_port_name) + 1 - ) # convert port name e.g. '80' to index - - found_action = True - break - - if found_action: - action = [ - action_decision, - action_permission, - action_source_id, - action_destination_id, - action_protocol, - action_port, - ] - action = transform_action_acl_enum(action) - action = get_new_action(action, self._env.action_dict) - else: - # If no good/useful action has been found, just perform a nothing action - action = ["NONE", "ALLOW", "ANY", "ANY", "ANY", "ANY"] - action = transform_action_acl_enum(action) - action = get_new_action(action, self._env.action_dict) - return action - - def _calculate_action_basic_view(self, obs: np.ndarray) -> int: - """ - Calculate a good acl-based action for the blue agent to take. - - Uses ONLY information from the current observation with NO knowledge - of previous actions taken and NO reward feedback. - - We rely on randomness to select the precise action, as we want to - block all traffic originating from a compromised node, without being - able to tell: - 1. Which ACL rules already exist - 2. Which actions the agent has already tried. - - There is a high probability that the correct rule will not be deleted - before the state becomes overwhelmed. - - Currently, a deny rule does not overwrite an allow rule. The allow - rules must be deleted. - - :param obs: current observation from the gym environment - :type obs: np.ndarray - :return: Optimal action to take in the environment (chosen from the discrete action space) - :rtype: int - """ - action_dict = self._env.action_dict - r_obs = transform_change_obs_readable(obs) - _, o, _, *s = r_obs - - if len(r_obs) == 4: # only 1 service - s = [*s] - - number_of_nodes = len([i for i in o if i != "NONE"]) # number of nodes (not links) - for service_num, service_states in enumerate(s): - comprimised_states = [n for n, i in enumerate(service_states) if i == "COMPROMISED"] - if len(comprimised_states) == 0: - # No states are COMPROMISED, try the next service - continue - - compromised_node = np.random.choice(comprimised_states) + 1 # +1 as 0 would be any - action_decision = "DELETE" - action_permission = "ALLOW" - action_source_ip = compromised_node - # Randomly select a destination ID to block - action_destination_ip = np.random.choice(list(range(1, number_of_nodes + 1)) + ["ANY"]) - action_destination_ip = ( - int(action_destination_ip) if action_destination_ip != "ANY" else action_destination_ip - ) - action_protocol = service_num + 1 # +1 as 0 is any - # Randomly select a port to block - # Bad assumption that number of protocols equals number of ports - # AND no rules exist with an ANY port - action_port = np.random.choice(list(range(1, len(s) + 1))) - - action = [ - action_decision, - action_permission, - action_source_ip, - action_destination_ip, - action_protocol, - action_port, - ] - action = transform_action_acl_enum(action) - action = get_new_action(action, action_dict) - # We can only perform 1 action on each step - return action - - # If no good/useful action has been found, just perform a nothing action - nothing_action = ["NONE", "ALLOW", "ANY", "ANY", "ANY", "ANY"] - nothing_action = transform_action_acl_enum(nothing_action) - nothing_action = get_new_action(nothing_action, action_dict) - return nothing_action diff --git a/src/primaite/agents/hardcoded_node.py b/src/primaite/agents/hardcoded_node.py deleted file mode 100644 index b08d8967..00000000 --- a/src/primaite/agents/hardcoded_node.py +++ /dev/null @@ -1,125 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import numpy as np - -from primaite.agents.hardcoded_abc import HardCodedAgentSessionABC -from primaite.agents.utils import get_new_action, transform_action_node_enum, transform_change_obs_readable - - -class HardCodedNodeAgent(HardCodedAgentSessionABC): - """An Agent Session class that implements a deterministic Node agent.""" - - def _calculate_action(self, obs: np.ndarray) -> int: - """ - Calculate a good node-based action for the blue agent to take. - - :param obs: current observation from the gym environment - :type obs: np.ndarray - :return: Optimal action to take in the environment (chosen from the discrete action space) - :rtype: int - """ - action_dict = self._env.action_dict - r_obs = transform_change_obs_readable(obs) - _, o, os, *s = r_obs - - if len(r_obs) == 4: # only 1 service - s = [*s] - - # Check in order of most important states (order doesn't currently - # matter, but it probably should) - # First see if any OS states are compromised - for x, os_state in enumerate(os): - if os_state == "COMPROMISED": - action_node_id = x + 1 - action_node_property = "OS" - property_action = "PATCHING" - action_service_index = 0 # does nothing isn't relevant for os - action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - action = transform_action_node_enum(action) - action = get_new_action(action, action_dict) - # We can only perform 1 action on each step - return action - - # Next, see if any Services are compromised - # We fix the compromised state before overwhelemd state, - # If a compromised entry node is fixed before the overwhelmed state is triggered, instruction is ignored - for service_num, service in enumerate(s): - for x, service_state in enumerate(service): - if service_state == "COMPROMISED": - action_node_id = x + 1 - action_node_property = "SERVICE" - property_action = "PATCHING" - action_service_index = service_num - - action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - action = transform_action_node_enum(action) - action = get_new_action(action, action_dict) - # We can only perform 1 action on each step - return action - - # Next, See if any services are overwhelmed - # perhaps this should be fixed automatically when the compromised PCs issues are also resolved - # Currently there's no reason that an Overwhelmed state cannot be resolved before resolving the compromised PCs - - for service_num, service in enumerate(s): - for x, service_state in enumerate(service): - if service_state == "OVERWHELMED": - action_node_id = x + 1 - action_node_property = "SERVICE" - property_action = "PATCHING" - action_service_index = service_num - - action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - action = transform_action_node_enum(action) - action = get_new_action(action, action_dict) - # We can only perform 1 action on each step - return action - - # Finally, turn on any off nodes - for x, operating_state in enumerate(o): - if os_state == "OFF": - action_node_id = x + 1 - action_node_property = "OPERATING" - property_action = "ON" # Why reset it when we can just turn it on - action_service_index = 0 # does nothing isn't relevant for operating state - action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - # TODO: transform_action_node_enum takes only one argument, not sure why two are given here. - action = transform_action_node_enum(action, action_dict) - action = get_new_action(action, action_dict) - # We can only perform 1 action on each step - return action - - # If no good actions, just go with an action that wont do any harm - action_node_id = 1 - action_node_property = "NONE" - property_action = "ON" - action_service_index = 0 - action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - action = transform_action_node_enum(action) - action = get_new_action(action, action_dict) - - return action diff --git a/src/primaite/agents/rllib.py b/src/primaite/agents/rllib.py deleted file mode 100644 index 96bb0737..00000000 --- a/src/primaite/agents/rllib.py +++ /dev/null @@ -1,287 +0,0 @@ -# # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# from __future__ import annotations - -# import json -# import shutil -# import zipfile -# from datetime import datetime -# from logging import Logger -# from pathlib import Path -# from typing import Any, Callable, Dict, Optional, Union -# from uuid import uuid4 - -# from primaite import getLogger -# from primaite.agents.agent_abc import AgentSessionABC -# from primaite.common.enums import AgentFramework, AgentIdentifier, SessionType -# from primaite.environment.primaite_env import Primaite - -# # from ray.rllib.algorithms import Algorithm -# # from ray.rllib.algorithms.a2c import A2CConfig -# # from ray.rllib.algorithms.ppo import PPOConfig -# # from ray.tune.logger import UnifiedLogger -# # from ray.tune.registry import register_env - - -# # from primaite.exceptions import RLlibAgentError - -# _LOGGER: Logger = getLogger(__name__) - - -# # TODO: verify type of env_config -# def _env_creator(env_config: Dict[str, Any]) -> Primaite: -# return Primaite( -# training_config_path=env_config["training_config_path"], -# lay_down_config_path=env_config["lay_down_config_path"], -# session_path=env_config["session_path"], -# timestamp_str=env_config["timestamp_str"], -# ) - -# # # TODO: verify type hint return type -# # def _custom_log_creator(session_path: Path) -> Callable[[Dict], UnifiedLogger]: -# # logdir = session_path / "ray_results" -# # logdir.mkdir(parents=True, exist_ok=True) - -# # def logger_creator(config: Dict) -> UnifiedLogger: -# # return UnifiedLogger(config, logdir, loggers=None) - -# return logger_creator - - -# # class RLlibAgent(AgentSessionABC): -# # """An AgentSession class that implements a Ray RLlib agent.""" - -# # def __init__( -# # self, -# # training_config_path: Optional[Union[str, Path]] = "", -# # lay_down_config_path: Optional[Union[str, Path]] = "", -# # session_path: Optional[Union[str, Path]] = None, -# # ) -> None: -# # """ -# # Initialise the RLLib Agent training session. - -# # :param training_config_path: YAML file containing configurable items defined in -# # `primaite.config.training_config.TrainingConfig` -# # :type training_config_path: Union[path, str] -# # :param lay_down_config_path: YAML file containing configurable items for generating network laydown. -# # :type lay_down_config_path: Union[path, str] -# # :raises ValueError: If the training config contains a bad value for agent_framework (should be "RLLIB") -# # :raises ValueError: If the training config contains a bad value for agent_identifies (should be `PPO` -# # or `A2C`) -# # """ -# # # TODO: implement RLlib agent loading -# # if session_path is not None: -# # msg = "RLlib agent loading has not been implemented yet" -# # _LOGGER.critical(msg) -# # raise NotImplementedError(msg) - -# # super().__init__(training_config_path, lay_down_config_path) -# # if self._training_config.session_type == SessionType.EVAL: -# # msg = "Cannot evaluate an RLlib agent that hasn't been through training yet." -# # _LOGGER.critical(msg) -# # raise RLlibAgentError(msg) -# # if not self._training_config.agent_framework == AgentFramework.RLLIB: -# # msg = f"Expected RLLIB agent_framework, " f"got {self._training_config.agent_framework}" -# # _LOGGER.error(msg) -# # raise ValueError(msg) -# # self._agent_config_class: Union[PPOConfig, A2CConfig] -# # if self._training_config.agent_identifier == AgentIdentifier.PPO: -# # self._agent_config_class = PPOConfig -# # elif self._training_config.agent_identifier == AgentIdentifier.A2C: -# # self._agent_config_class = A2CConfig -# # else: -# # msg = "Expected PPO or A2C agent_identifier, " f"got {self._training_config.agent_identifier.value}" -# # _LOGGER.error(msg) -# # raise ValueError(msg) -# # self._agent_config: Union[PPOConfig, A2CConfig] - -# # self._current_result: dict -# # self._setup() -# # _LOGGER.debug( -# # f"Created {self.__class__.__name__} using: " -# # f"agent_framework={self._training_config.agent_framework}, " -# # f"agent_identifier=" -# # f"{self._training_config.agent_identifier}, " -# # f"deep_learning_framework=" -# # f"{self._training_config.deep_learning_framework}" -# # ) -# # self._train_agent = None # Required to capture the learning agent to close after eval - -# # def _update_session_metadata_file(self) -> None: -# # """ -# # Update the ``session_metadata.json`` file. - -# # Updates the `session_metadata.json`` in the ``session_path`` directory -# # with the following key/value pairs: - -# # - end_datetime: The date & time the session ended in iso format. -# # - total_episodes: The total number of training episodes completed. -# # - total_time_steps: The total number of training time steps completed. -# # """ -# # with open(self.session_path / "session_metadata.json", "r") as file: -# # metadata_dict = json.load(file) - -# # metadata_dict["end_datetime"] = datetime.now().isoformat() -# # if not self.is_eval: -# # metadata_dict["learning"]["total_episodes"] = self._current_result["episodes_total"] # noqa -# # metadata_dict["learning"]["total_time_steps"] = self._current_result["timesteps_total"] # noqa -# # else: -# # metadata_dict["evaluation"]["total_episodes"] = self._current_result["episodes_total"] # noqa -# # metadata_dict["evaluation"]["total_time_steps"] = self._current_result["timesteps_total"] # noqa - -# # filepath = self.session_path / "session_metadata.json" -# # _LOGGER.debug(f"Updating Session Metadata file: {filepath}") -# # with open(filepath, "w") as file: -# # json.dump(metadata_dict, file) -# # _LOGGER.debug("Finished updating session metadata file") - -# # def _setup(self) -> None: -# # super()._setup() -# # register_env("primaite", _env_creator) -# # self._agent_config = self._agent_config_class() - -# # self._agent_config.environment( -# # env="primaite", -# # env_config=dict( -# # training_config_path=self._training_config_path, -# # lay_down_config_path=self._lay_down_config_path, -# # session_path=self.session_path, -# # timestamp_str=self.timestamp_str, -# # ), -# # ) -# # self._agent_config.seed = self._training_config.seed - -# # self._agent_config.training(train_batch_size=self._training_config.num_train_steps) -# # self._agent_config.framework(framework="tf") - -# # self._agent_config.rollouts( -# # num_rollout_workers=1, -# # num_envs_per_worker=1, -# # horizon=self._training_config.num_train_steps, -# # ) -# # self._agent: Algorithm = self._agent_config.build(logger_creator=_custom_log_creator(self.learning_path)) - -# # def _save_checkpoint(self) -> None: -# # checkpoint_n = self._training_config.checkpoint_every_n_episodes -# # episode_count = self._current_result["episodes_total"] -# # save_checkpoint = False -# # if checkpoint_n: -# # save_checkpoint = episode_count % checkpoint_n == 0 -# # if episode_count and save_checkpoint: -# # self._agent.save(str(self.checkpoints_path)) - -# # def learn( -# # self, -# # **kwargs: Any, -# # ) -> None: -# # """ -# # Evaluate the agent. - -# # :param kwargs: Any agent-specific key-word args to be passed. -# # """ -# # time_steps = self._training_config.num_train_steps -# # episodes = self._training_config.num_train_episodes - -# # _LOGGER.info(f"Beginning learning for {episodes} episodes @" f" {time_steps} time steps...") -# # for i in range(episodes): -# # self._current_result = self._agent.train() -# # self._save_checkpoint() -# # self.save() -# # super().learn() -# # # Done this way as the RLlib eval can only be performed if the session hasn't been stopped -# # if self._training_config.session_type is not SessionType.TRAIN: -# # self._train_agent = self._agent -# # else: -# # self._agent.stop() -# # self._plot_av_reward_per_episode(learning_session=True) - -# # def _unpack_saved_agent_into_eval(self) -> Path: -# # """Unpacks the pre-trained and saved RLlib agent so that it can be reloaded by Ray for eval.""" -# # agent_restore_path = self.evaluation_path / "agent_restore" -# # if agent_restore_path.exists(): -# # shutil.rmtree(agent_restore_path) -# # agent_restore_path.mkdir() -# # with zipfile.ZipFile(self._saved_agent_path, "r") as zip_file: -# # zip_file.extractall(agent_restore_path) -# # return agent_restore_path - -# # def _setup_eval(self): -# # self._can_learn = False -# # self._can_evaluate = True -# # self._agent.restore(str(self._unpack_saved_agent_into_eval())) - -# # def evaluate( -# # self, -# # **kwargs, -# # ): -# # """ -# # Evaluate the agent. - -# # :param kwargs: Any agent-specific key-word args to be passed. -# # """ -# # time_steps = self._training_config.num_eval_steps -# # episodes = self._training_config.num_eval_episodes - -# # self._setup_eval() - -# # self._env: Primaite = Primaite( -# # self._training_config_path, self._lay_down_config_path, self.session_path, self.timestamp_str -# # ) - -# # self._env.set_as_eval() -# # self.is_eval = True -# # if self._training_config.deterministic: -# # deterministic_str = "deterministic" -# # else: -# # deterministic_str = "non-deterministic" -# # _LOGGER.info( -# # f"Beginning {deterministic_str} evaluation for " f"{episodes} episodes @ {time_steps} time steps..." -# # ) -# # for episode in range(episodes): -# # obs = self._env.reset() -# # for step in range(time_steps): -# # action = self._agent.compute_single_action(observation=obs, explore=False) - -# # obs, rewards, done, info = self._env.step(action) - -# # self._env.reset() -# # self._env.close() -# # super().evaluate() -# # # Now we're safe to close the learning agent and write the mean rewards per episode for it -# # if self._training_config.session_type is not SessionType.TRAIN: -# # self._train_agent.stop() -# # self._plot_av_reward_per_episode(learning_session=True) -# # # Perform a clean-up of the unpacked agent -# # if (self.evaluation_path / "agent_restore").exists(): -# # shutil.rmtree((self.evaluation_path / "agent_restore")) - -# # def _get_latest_checkpoint(self) -> None: -# # raise NotImplementedError - -# # @classmethod -# # def load(cls, path: Union[str, Path]) -> RLlibAgent: -# # """Load an agent from file.""" -# # raise NotImplementedError - -# # def save(self, overwrite_existing: bool = True) -> None: -# # """Save the agent.""" -# # # Make temp dir to save in isolation -# # temp_dir = self.learning_path / str(uuid4()) -# # temp_dir.mkdir() - -# # # Save the agent to the temp dir -# # self._agent.save(str(temp_dir)) - -# # # Capture the saved Rllib checkpoint inside the temp directory -# # for file in temp_dir.iterdir(): -# # checkpoint_dir = file -# # break - -# # # Zip the folder -# # shutil.make_archive(str(self._saved_agent_path).replace(".zip", ""), "zip", checkpoint_dir) # noqa - -# # # Drop the temp directory -# # shutil.rmtree(temp_dir) - -# # def export(self) -> None: -# # """Export the agent to transportable file format.""" -# # raise NotImplementedError diff --git a/src/primaite/agents/sb3.py b/src/primaite/agents/sb3.py deleted file mode 100644 index 92c5ee5f..00000000 --- a/src/primaite/agents/sb3.py +++ /dev/null @@ -1,206 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from __future__ import annotations - -import json -from logging import Logger -from pathlib import Path -from typing import Any, Optional, Union - -import numpy as np -from stable_baselines3 import A2C, PPO -from stable_baselines3.ppo import MlpPolicy as PPOMlp - -from primaite import getLogger -from primaite.agents.agent_abc import AgentSessionABC -from primaite.common.enums import AgentFramework, AgentIdentifier -from primaite.environment.primaite_env import Primaite - -_LOGGER: Logger = getLogger(__name__) - - -class SB3Agent(AgentSessionABC): - """An AgentSession class that implements a Stable Baselines3 agent.""" - - def __init__( - self, - training_config_path: Optional[Union[str, Path]] = None, - lay_down_config_path: Optional[Union[str, Path]] = None, - session_path: Optional[Union[str, Path]] = None, - legacy_training_config: bool = False, - legacy_lay_down_config: bool = False, - ) -> None: - """ - Initialise the SB3 Agent training session. - - :param training_config_path: YAML file containing configurable items defined in - `primaite.config.training_config.TrainingConfig` - :type training_config_path: Union[path, str] - :param lay_down_config_path: YAML file containing configurable items for generating network laydown. - :type lay_down_config_path: Union[path, str] - :param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :raises ValueError: If the training config contains an unexpected value for agent_framework (should be "SB3") - :raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO` - or `A2C`) - """ - super().__init__( - training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config - ) - if not self._training_config.agent_framework == AgentFramework.SB3: - msg = f"Expected SB3 agent_framework, " f"got {self._training_config.agent_framework}" - _LOGGER.error(msg) - raise ValueError(msg) - self._agent_class: Union[PPO, A2C] - if self._training_config.agent_identifier == AgentIdentifier.PPO: - self._agent_class = PPO - elif self._training_config.agent_identifier == AgentIdentifier.A2C: - self._agent_class = A2C - else: - msg = "Expected PPO or A2C agent_identifier, " f"got {self._training_config.agent_identifier}" - _LOGGER.error(msg) - raise ValueError(msg) - - self._tensorboard_log_path = self.learning_path / "tensorboard_logs" - self._tensorboard_log_path.mkdir(parents=True, exist_ok=True) - - _LOGGER.debug( - f"Created {self.__class__.__name__} using: " - f"agent_framework={self._training_config.agent_framework}, " - f"agent_identifier=" - f"{self._training_config.agent_identifier}" - ) - - self.is_eval = False - - self._setup() - - def _setup(self) -> None: - """Set up the SB3 Agent.""" - self._env = Primaite( - training_config_path=self._training_config_path, - lay_down_config_path=self._lay_down_config_path, - session_path=self.session_path, - timestamp_str=self.timestamp_str, - legacy_training_config=self.legacy_training_config, - legacy_lay_down_config=self.legacy_lay_down_config, - ) - - # check if there is a zip file that needs to be loaded - load_file = next(self.session_path.rglob("*.zip"), None) - - if not load_file: - # create a new env and agent - - self._agent = self._agent_class( - PPOMlp, - self._env, - verbose=self.sb3_output_verbose_level, - n_steps=self._training_config.num_train_steps, - tensorboard_log=str(self._tensorboard_log_path), - seed=self._training_config.seed, - ) - else: - # set env values from session metadata - with open(self.session_path / "session_metadata.json", "r") as file: - md_dict = json.load(file) - - # load environment values - if self.is_eval: - # evaluation always starts at 0 - self._env.episode_count = 0 - self._env.total_step_count = 0 - else: - # carry on from previous learning sessions - self._env.episode_count = md_dict["learning"]["total_episodes"] - self._env.total_step_count = md_dict["learning"]["total_time_steps"] - - # load the file - self._agent = self._agent_class.load(load_file, env=self._env) - - # set agent values - self._agent.verbose = self.sb3_output_verbose_level - self._agent.tensorboard_log = self.session_path / "learning/tensorboard_logs" - - super()._setup() - - def _save_checkpoint(self) -> None: - checkpoint_n = self._training_config.checkpoint_every_n_episodes - episode_count = self._env.episode_count - save_checkpoint = False - if checkpoint_n: - save_checkpoint = episode_count % checkpoint_n == 0 - if episode_count and save_checkpoint: - checkpoint_path = self.checkpoints_path / f"sb3ppo_{episode_count}.zip" - self._agent.save(checkpoint_path) - _LOGGER.debug(f"Saved agent checkpoint: {checkpoint_path}") - - def _get_latest_checkpoint(self) -> None: - pass - - def learn( - self, - **kwargs: Any, - ) -> None: - """ - Train the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - time_steps = self._training_config.num_train_steps - episodes = self._training_config.num_train_episodes - self.is_eval = False - _LOGGER.info(f"Beginning learning for {episodes} episodes @" f" {time_steps} time steps...") - for i in range(episodes): - self._agent.learn(total_timesteps=time_steps) - self._save_checkpoint() - self._env._write_av_reward_per_episode() # noqa - self.save() - self._env.close() - super().learn() - - # save agent - self.save() - - self._plot_av_reward_per_episode(learning_session=True) - - def evaluate( - self, - **kwargs: Any, - ) -> None: - """ - Evaluate the agent. - - :param kwargs: Any agent-specific key-word args to be passed. - """ - time_steps = self._training_config.num_eval_steps - episodes = self._training_config.num_eval_episodes - self._env.set_as_eval() - self.is_eval = True - if self._training_config.deterministic: - deterministic_str = "deterministic" - else: - deterministic_str = "non-deterministic" - _LOGGER.info( - f"Beginning {deterministic_str} evaluation for " f"{episodes} episodes @ {time_steps} time steps..." - ) - for episode in range(episodes): - obs = self._env.reset() - - for step in range(time_steps): - action, _states = self._agent.predict(obs, deterministic=self._training_config.deterministic) - if isinstance(action, np.ndarray): - action = np.int64(action) - obs, rewards, done, info = self._env.step(action) - self._env._write_av_reward_per_episode() # noqa - self._env.close() - super().evaluate() - - def save(self) -> None: - """Save the agent.""" - self._agent.save(self._saved_agent_path) - - def export(self) -> None: - """Export the agent to transportable file format.""" - raise NotImplementedError diff --git a/src/primaite/agents/simple.py b/src/primaite/agents/simple.py deleted file mode 100644 index bfdff869..00000000 --- a/src/primaite/agents/simple.py +++ /dev/null @@ -1,59 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -import numpy as np - -from primaite.agents.hardcoded_abc import HardCodedAgentSessionABC -from primaite.agents.utils import get_new_action, transform_action_acl_enum, transform_action_node_enum - - -class RandomAgent(HardCodedAgentSessionABC): - """ - A Random Agent. - - Get a completely random action from the action space. - """ - - def _calculate_action(self, obs: np.ndarray) -> int: - return self._env.action_space.sample() - - -class DummyAgent(HardCodedAgentSessionABC): - """ - A Dummy Agent. - - All action spaces setup so dummy action is always 0 regardless of action type used. - """ - - def _calculate_action(self, obs: np.ndarray) -> int: - return 0 - - -class DoNothingACLAgent(HardCodedAgentSessionABC): - """ - A do nothing ACL agent. - - A valid ACL action that has no effect; does nothing. - """ - - def _calculate_action(self, obs: np.ndarray) -> int: - nothing_action = ["NONE", "ALLOW", "ANY", "ANY", "ANY", "ANY"] - nothing_action = transform_action_acl_enum(nothing_action) - nothing_action = get_new_action(nothing_action, self._env.action_dict) - - return nothing_action - - -class DoNothingNodeAgent(HardCodedAgentSessionABC): - """ - A do nothing Node agent. - - A valid Node action that has no effect; does nothing. - """ - - def _calculate_action(self, obs: np.ndarray) -> int: - nothing_action = [1, "NONE", "ON", 0] - nothing_action = transform_action_node_enum(nothing_action) - nothing_action = get_new_action(nothing_action, self._env.action_dict) - # nothing_action should currently always be 0 - - return nothing_action diff --git a/src/primaite/agents/utils.py b/src/primaite/agents/utils.py deleted file mode 100644 index 08d46294..00000000 --- a/src/primaite/agents/utils.py +++ /dev/null @@ -1,450 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from typing import Dict, List, Union - -import numpy as np - -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import ( - HardwareState, - LinkStatus, - NodeHardwareAction, - NodePOLType, - NodeSoftwareAction, - SoftwareState, -) - - -def transform_action_node_readable(action: List[int]) -> List[Union[int, str]]: - """Convert a node action from enumerated format to readable format. - - example: - [1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0] - - :param action: Agent action, formatted as a list of ints, for more information check out - `primaite.environment.primaite_env.Primaite` - :type action: List[int] - :return: The same action list, but with the encodings translated back into meaningful labels - :rtype: List[Union[int,str]] - """ - action_node_property = NodePOLType(action[1]).name - - if action_node_property == "OPERATING": - property_action = NodeHardwareAction(action[2]).name - elif (action_node_property == "OS" or action_node_property == "SERVICE") and action[2] <= 1: - property_action = NodeSoftwareAction(action[2]).name - else: - property_action = "NONE" - - new_action: list[Union[int, str]] = [action[0], action_node_property, property_action, action[3]] - return new_action - - -def transform_action_acl_readable(action: List[int]) -> List[Union[str, int]]: - """ - Transform an ACL action to a more readable format. - - example: - [0, 1, 2, 5, 0, 1] -> ['NONE', 'ALLOW', 2, 5, 'ANY', 1] - - :param action: Agent action, formatted as a list of ints, for more information check out - `primaite.environment.primaite_env.Primaite` - :type action: List[int] - :return: The same action list, but with the encodings translated back into meaningful labels - :rtype: List[Union[int,str]] - """ - action_decisions = {0: "NONE", 1: "CREATE", 2: "DELETE"} - action_permissions = {0: "DENY", 1: "ALLOW"} - - action_decision = action_decisions[action[0]] - action_permission = action_permissions[action[1]] - - # For IPs, Ports and Protocols, 0 means any, otherwise its just an index - new_action = [action_decision, action_permission] + list(action[2:6]) - for n, val in enumerate(list(action[2:6])): - if val == 0: - new_action[n + 2] = "ANY" - - return new_action - - -def is_valid_node_action(action: List[int]) -> bool: - """ - Is the node action an actual valid action. - - Only uses information about the action to determine if the action has an effect - - Does NOT consider: - - Node ID not valid to perform an operation - e.g. selected node has no service so cannot patch - - Node already being in that state (turning an ON node ON) - - :param action: Agent action, formatted as a list of ints, for more information check out - `primaite.environment.primaite_env.Primaite` - :type action: List[int] - :return: Whether the action is valid - :rtype: bool - """ - action_r = transform_action_node_readable(action) - - node_property = action_r[1] - node_action = action_r[2] - - # print("node property", node_property, "\nnode action", node_action) - - if node_property == "NONE": - return False - if node_action == "NONE": - return False - if node_property == "OPERATING" and node_action == "PATCHING": - # Operating State cannot PATCH - return False - if node_property != "OPERATING" and node_action not in [ - "NONE", - "PATCHING", - ]: - # Software States can only do Nothing or Patch - return False - return True - - -def is_valid_acl_action(action: List[int]) -> bool: - """ - Is the ACL action an actual valid action. - - Only uses information about the action to determine if the action has an effect. - - Does NOT consider: - - Trying to create identical rules - - Trying to create a rule which is a subset of another rule (caused by "ANY") - - :param action: Agent action, formatted as a list of ints, for more information check out - `primaite.environment.primaite_env.Primaite` - :type action: List[int] - :return: Whether the action is valid - :rtype: bool - """ - action_r = transform_action_acl_readable(action) - - action_decision = action_r[0] - action_permission = action_r[1] - action_source_id = action_r[2] - action_destination_id = action_r[3] - - if action_decision == "NONE": - return False - if action_source_id == action_destination_id and action_source_id != "ANY" and action_destination_id != "ANY": - # ACL rule towards itself - return False - if action_permission == "DENY": - # DENY is unnecessary, we can create and delete allow rules instead - # No allow rule = blocked/DENY by feault. ALLOW overrides existing DENY. - return False - - return True - - -def is_valid_acl_action_extra(action: List[int]) -> bool: - """ - Harsher version of valid acl actions, does not allow action. - - :param action: Agent action, formatted as a list of ints, for more information check out - `primaite.environment.primaite_env.Primaite` - :type action: List[int] - :return: Whether the action is valid - :rtype: bool - """ - if is_valid_acl_action(action) is False: - return False - - action_r = transform_action_acl_readable(action) - action_protocol = action_r[4] - action_port = action_r[5] - - # Don't allow protocols or ports to be ANY - # in the future we might want to do the opposite, and only have ANY option for ports and service - if action_protocol == "ANY": - return False - if action_port == "ANY": - return False - - return True - - -def transform_change_obs_readable(obs: np.ndarray) -> List[List[Union[str, int]]]: - """Transform list of transactions to readable list of each observation property. - - example: - np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 2], ['OFF', 'ON'], ['GOOD', 'GOOD'], ['COMPROMISED', 'GOOD']] - - :param obs: Raw observation from the environment. - :type obs: np.ndarray - :return: The same observation, but the encoded integer values are replaced with readable names. - :rtype: List[List[Union[str, int]]] - """ - ids = [i for i in obs[:, 0]] - operating_states = [HardwareState(i).name for i in obs[:, 1]] - os_states = [SoftwareState(i).name for i in obs[:, 2]] - new_obs = [ids, operating_states, os_states] - - for service in range(4, obs.shape[1]): - # Links bit/s don't have a service state - service_states = [SoftwareState(i).name if i <= 4 else i for i in obs[:, service]] - new_obs.append(service_states) - - return new_obs - - -def transform_obs_readable(obs: np.ndarray) -> List[List[Union[str, int]]]: - """Transform observation to readable format. - - example - np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 'OFF', 'GOOD', 'COMPROMISED'], [2, 'ON', 'GOOD', 'GOOD']] - - :param obs: Raw observation from the environment. - :type obs: np.ndarray - :return: The same observation, but the encoded integer values are replaced with readable names. - :rtype: List[List[Union[str, int]]] - """ - changed_obs = transform_change_obs_readable(obs) - new_obs = list(zip(*changed_obs)) - # Convert list of tuples to list of lists - new_obs = [list(i) for i in new_obs] - - return new_obs - - -def convert_to_new_obs(obs: np.ndarray, num_nodes: int = 10) -> np.ndarray: - """Convert original gym Box observation space to new multiDiscrete observation space. - - :param obs: observation in the 'old' (NodeLinkTable) format - :type obs: np.ndarray - :param num_nodes: number of nodes in the network, defaults to 10 - :type num_nodes: int, optional - :return: reformatted observation - :rtype: np.ndarray - """ - # Remove ID columns, remove links and flatten to MultiDiscrete observation space - new_obs = obs[:num_nodes, 1:].flatten() - return new_obs - - -def convert_to_old_obs(obs: np.ndarray, num_nodes: int = 10, num_links: int = 10, num_services: int = 1) -> np.ndarray: - """Convert to old observation. - - Links filled with 0's as no information is included in new observation space. - - example: - obs = array([1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1, 1, 1]) - - new_obs = array([[ 1, 1, 1, 1], - [ 2, 1, 1, 1], - [ 3, 1, 1, 1], - ... - [20, 0, 0, 0]]) - - :param obs: observation in the 'new' (MultiDiscrete) format - :type obs: np.ndarray - :param num_nodes: number of nodes in the network, defaults to 10 - :type num_nodes: int, optional - :param num_links: number of links in the network, defaults to 10 - :type num_links: int, optional - :param num_services: number of services on the network, defaults to 1 - :type num_services: int, optional - :return: 2-d BOX observation space, in the same format as NodeLinkTable - :rtype: np.ndarray - """ - # Convert back to more readable, original format - reshaped_nodes = obs[:-num_links].reshape(num_nodes, num_services + 2) - - # Add empty links back and add node ID back - s = np.zeros( - [reshaped_nodes.shape[0] + num_links, reshaped_nodes.shape[1] + 1], - dtype=np.int64, - ) - s[:, 0] = range(1, num_nodes + num_links + 1) # Adding ID back - s[:num_nodes, 1:] = reshaped_nodes # put values back in - new_obs = s - - # Add links back in - links = obs[-num_links:] - # Links will be added to the last protocol/service slot but they are not specific to that service - new_obs[num_nodes:, -1] = links - - return new_obs - - -def describe_obs_change( - obs1: np.ndarray, obs2: np.ndarray, num_nodes: int = 10, num_links: int = 10, num_services: int = 1 -) -> str: - """Build a string describing the difference between two observations. - - example: - obs_1 = array([[1, 1, 1, 1, 3], [2, 1, 1, 1, 1]]) - obs_2 = array([[1, 1, 1, 1, 1], [2, 1, 1, 1, 1]]) - output = 'ID 1: SERVICE 2 set to GOOD' - - :param obs1: First observation - :type obs1: np.ndarray - :param obs2: Second observation - :type obs2: np.ndarray - :param num_nodes: How many nodes are in the network laydown, defaults to 10 - :type num_nodes: int, optional - :param num_links: How many links are in the network laydown, defaults to 10 - :type num_links: int, optional - :param num_services: How many services are configured for this scenario, defaults to 1 - :type num_services: int, optional - :return: A multi-line string with a human-readable description of the difference. - :rtype: str - """ - obs1 = convert_to_old_obs(obs1, num_nodes, num_links, num_services) - obs2 = convert_to_old_obs(obs2, num_nodes, num_links, num_services) - list_of_changes = [] - for n, row in enumerate(obs1 - obs2): - if row.any() != 0: - relevant_changes = np.where(row != 0, obs2[n], -1) - relevant_changes[0] = obs2[n, 0] # ID is always relevant - is_link = relevant_changes[0] > num_nodes - desc = _describe_obs_change_helper(relevant_changes, is_link) - list_of_changes.append(desc) - - change_string = "\n ".join(list_of_changes) - if len(list_of_changes) > 0: - change_string = "\n " + change_string - return change_string - - -def _describe_obs_change_helper(obs_change: List[int], is_link: bool) -> str: - """ - Helper funcion to describe what has changed. - - example: - [ 1 -1 -1 -1 1] -> "ID 1: Service 1 changed to GOOD" - - Handles multiple changes e.g. 'ID 1: SERVICE 1 changed to PATCHING. SERVICE 2 set to GOOD.' - - :param obs_change: List of integers generated within the `describe_obs_change` function. It should correspond to one - row of the observation table, and have `-1` at locations where the observation hasn't changed, and the new - status where it has changed. - :type obs_change: List[int] - :param is_link: Whether the row of the observation space corresponds to a link. False means it represents a node. - :type is_link: bool - :return: A human-readable description of the difference between the two observation rows. - :rtype: str - """ - # Indexes where a change has occured, not including 0th index - index_changed = [i for i in range(1, len(obs_change)) if obs_change[i] != -1] - # Node pol types, Indexes >= 3 are service nodes - NodePOLTypes = [NodePOLType(i).name if i < 3 else NodePOLType(3).name + " " + str(i - 3) for i in index_changed] - # Account for hardware states, software sattes and links - states = [ - LinkStatus(obs_change[i]).name - if is_link - else HardwareState(obs_change[i]).name - if i == 1 - else SoftwareState(obs_change[i]).name - for i in index_changed - ] - - if not is_link: - desc = f"ID {obs_change[0]}:" - for node_pol_type, state in list(zip(NodePOLTypes, states)): - desc = desc + " " + node_pol_type + " changed to " + state + "." - else: - desc = f"ID {obs_change[0]}: Link traffic changed to {states[0]}." - - return desc - - -def transform_action_node_enum(action: List[Union[str, int]]) -> List[int]: - """Convert a node action from readable string format, to enumerated format. - - example: - [1, 'SERVICE', 'PATCHING', 0] -> [1, 3, 1, 0] - :param action: Action in 'readable' format - :type action: List[Union[str,int]] - :return: Action with verbs encoded as ints - :rtype: List[int] - """ - action_node_id = action[0] - action_node_property = NodePOLType[action[1]].value - - if action[1] == "OPERATING": - property_action = NodeHardwareAction[action[2]].value - elif action[1] == "OS" or action[1] == "SERVICE": - property_action = NodeSoftwareAction[action[2]].value - else: - property_action = 0 - - action_service_index = action[3] - - new_action = [ - action_node_id, - action_node_property, - property_action, - action_service_index, - ] - - return new_action - - -def transform_action_acl_enum(action: List[Union[int, str]]) -> np.ndarray: - """ - Convert acl action from readable str format, to enumerated format. - - :param action: ACL-based action expressed as a list of human-readable ints and strings - :type action: List[Union[int,str]] - :return: The same action but encoded to contain only integers. - :rtype: np.ndarray - """ - action_decisions = {"NONE": 0, "CREATE": 1, "DELETE": 2} - action_permissions = {"DENY": 0, "ALLOW": 1} - - action_decision = action_decisions[action[0]] - action_permission = action_permissions[action[1]] - - # For IPs, Ports and Protocols, ANY has value 0, otherwise its just an index - new_action = [action_decision, action_permission] + list(action[2:6]) - for n, val in enumerate(list(action[2:6])): - if val == "ANY": - new_action[n + 2] = 0 - - new_action = np.array(new_action) - return new_action - - -def get_node_of_ip(ip: str, node_dict: Dict[str, NodeUnion]) -> str: - """Get the node ID of an IP address. - - node_dict: dictionary of nodes where key is ID, and value is the node (can be ontained from env.nodes) - - :param ip: The IP address of the node whose ID is required - :type ip: str - :param node_dict: The environment's node registry dictionary - :type node_dict: Dict[str,NodeUnion] - :return: The key from the registry dict that corresponds to the node with the IP adress provided by `ip` - :rtype: str - """ - for node_key, node_value in node_dict.items(): - node_ip = node_value.ip_address - if node_ip == ip: - return node_key - - -def get_new_action(old_action: np.ndarray, action_dict: Dict[int, List]) -> int: - """ - Get new action (e.g. 32) from old action e.g. [1,1,1,0]. - - Old_action can be either node or acl action type - - :param old_action: Action expressed as a list of choices, eg. [1,1,1,0] - :type old_action: np.ndarray - :param action_dict: Dictionary for translating the multidiscrete actions into the list-based actions. - :type action_dict: Dict[int,List] - :return: Action key correspoinding to the input `old_action` - :rtype: int - """ - for key, val in action_dict.items(): - if list(val) == list(old_action): - return key - # Not all possible actions are included in dict, only valid action are - # if action is not in the dict, its an invalid action so return 0 - return 0 diff --git a/src/primaite/common/__init__.py b/src/primaite/common/__init__.py deleted file mode 100644 index 5770bcbc..00000000 --- a/src/primaite/common/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Objects which are shared between many PrimAITE modules.""" diff --git a/src/primaite/common/custom_typing.py b/src/primaite/common/custom_typing.py deleted file mode 100644 index 4130e71a..00000000 --- a/src/primaite/common/custom_typing.py +++ /dev/null @@ -1,8 +0,0 @@ -from typing import Union - -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.passive_node import PassiveNode -from primaite.nodes.service_node import ServiceNode - -NodeUnion = Union[ActiveNode, PassiveNode, ServiceNode] -"""A Union of ActiveNode, PassiveNode, and ServiceNode.""" diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py deleted file mode 100644 index c33e764b..00000000 --- a/src/primaite/common/enums.py +++ /dev/null @@ -1,208 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Enumerations for APE.""" - -from enum import Enum, IntEnum - - -class NodeType(Enum): - """Node type enumeration.""" - - CCTV = 1 - SWITCH = 2 - COMPUTER = 3 - LINK = 4 - MONITOR = 5 - PRINTER = 6 - LOP = 7 - RTU = 8 - ACTUATOR = 9 - SERVER = 10 - - -class Priority(Enum): - """Node priority enumeration.""" - - P1 = 1 - P2 = 2 - P3 = 3 - P4 = 4 - P5 = 5 - - -class HardwareState(Enum): - """Node hardware state enumeration.""" - - NONE = 0 - ON = 1 - OFF = 2 - RESETTING = 3 - SHUTTING_DOWN = 4 - BOOTING = 5 - - -class SoftwareState(Enum): - """Software or Service state enumeration.""" - - NONE = 0 - GOOD = 1 - PATCHING = 2 - COMPROMISED = 3 - OVERWHELMED = 4 - - -class NodePOLType(Enum): - """Node Pattern of Life type enumeration.""" - - NONE = 0 - OPERATING = 1 - OS = 2 - SERVICE = 3 - FILE = 4 - - -class NodePOLInitiator(Enum): - """Node Pattern of Life initiator enumeration.""" - - DIRECT = 1 - IER = 2 - SERVICE = 3 - - -class Protocol(Enum): - """Service protocol enumeration.""" - - LDAP = 0 - FTP = 1 - HTTPS = 2 - SMTP = 3 - RTP = 4 - IPP = 5 - TCP = 6 - NONE = 7 - - -class SessionType(Enum): - """The type of PrimAITE Session to be run.""" - - TRAIN = 1 - "Train an agent" - EVAL = 2 - "Evaluate an agent" - TRAIN_EVAL = 3 - "Train then evaluate an agent" - - -class AgentFramework(Enum): - """The agent algorithm framework/package.""" - - CUSTOM = 0 - "Custom Agent" - SB3 = 1 - "Stable Baselines3" - # RLLIB = 2 - # "Ray RLlib" - - -class DeepLearningFramework(Enum): - """The deep learning framework.""" - - TF = "tf" - "Tensorflow" - TF2 = "tf2" - "Tensorflow 2.x" - TORCH = "torch" - "PyTorch" - - -class AgentIdentifier(Enum): - """The Red Agent algo/class.""" - - A2C = 1 - "Advantage Actor Critic" - PPO = 2 - "Proximal Policy Optimization" - HARDCODED = 3 - "The Hardcoded agents" - DO_NOTHING = 4 - "The DoNothing agents" - RANDOM = 5 - "The RandomAgent" - DUMMY = 6 - "The DummyAgent" - - -class HardCodedAgentView(Enum): - """The view the deterministic hard-coded agent has of the environment.""" - - BASIC = 1 - "The current observation space only" - FULL = 2 - "Full environment view with actions taken and reward feedback" - - -class ActionType(Enum): - """Action type enumeration.""" - - NODE = 0 - ACL = 1 - ANY = 2 - - -# TODO: this is not used anymore, write a ticket to delete it. -class ObservationType(Enum): - """Observation type enumeration.""" - - BOX = 0 - MULTIDISCRETE = 1 - - -class FileSystemState(Enum): - """File System State.""" - - GOOD = 1 - CORRUPT = 2 - DESTROYED = 3 - REPAIRING = 4 - RESTORING = 5 - - -class NodeHardwareAction(Enum): - """Node hardware action.""" - - NONE = 0 - ON = 1 - OFF = 2 - RESET = 3 - - -class NodeSoftwareAction(Enum): - """Node software action.""" - - NONE = 0 - PATCHING = 1 - - -class LinkStatus(Enum): - """Link traffic status.""" - - NONE = 0 - LOW = 1 - MEDIUM = 2 - HIGH = 3 - OVERLOAD = 4 - - -class SB3OutputVerboseLevel(IntEnum): - """The Stable Baselines3 learn/eval output verbosity level.""" - - NONE = 0 - INFO = 1 - DEBUG = 2 - - -class RulePermissionType(Enum): - """Any firewall rule type.""" - - NONE = 0 - DENY = 1 - ALLOW = 2 diff --git a/src/primaite/common/protocol.py b/src/primaite/common/protocol.py deleted file mode 100644 index 6940ba3f..00000000 --- a/src/primaite/common/protocol.py +++ /dev/null @@ -1,47 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The protocol class.""" - - -class Protocol(object): - """Protocol class.""" - - def __init__(self, _name: str) -> None: - """ - Initialise a protocol. - - :param _name: The name of the protocol - :type _name: str - """ - self.name: str = _name - self.load: int = 0 # bps - - def get_name(self) -> str: - """ - Gets the protocol name. - - Returns: - The protocol name - """ - return self.name - - def get_load(self) -> int: - """ - Gets the protocol load. - - Returns: - The protocol load (bps) - """ - return self.load - - def add_load(self, _load: int) -> None: - """ - Adds load to the protocol. - - Args: - _load: The load to add - """ - self.load += _load - - def clear_load(self) -> None: - """Clears the load on this protocol.""" - self.load = 0 diff --git a/src/primaite/common/service.py b/src/primaite/common/service.py deleted file mode 100644 index 956815e8..00000000 --- a/src/primaite/common/service.py +++ /dev/null @@ -1,28 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The Service class.""" - -from primaite.common.enums import SoftwareState - - -class Service(object): - """Service class.""" - - def __init__(self, name: str, port: str, software_state: SoftwareState) -> None: - """ - Initialise a service. - - :param name: The service name. - :param port: The service port. - :param software_state: The service SoftwareState. - """ - self.name: str = name - self.port: str = port - self.software_state: SoftwareState = software_state - self.patching_count: int = 0 - - def reduce_patching_count(self) -> None: - """Reduces the patching count for the service.""" - self.patching_count -= 1 - if self.patching_count <= 0: - self.patching_count = 0 - self.software_state = SoftwareState.GOOD diff --git a/src/primaite/config/__init__.py b/src/primaite/config/__init__.py deleted file mode 100644 index 92f5a7d2..00000000 --- a/src/primaite/config/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Configuration parameters for running experiments.""" diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml deleted file mode 100644 index dad0ff4b..00000000 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml +++ /dev/null @@ -1,166 +0,0 @@ -- item_type: PORTS - ports_list: - - port: '80' -- item_type: SERVICES - service_list: - - name: TCP -- item_type: NODE - node_id: '1' - name: PC1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.2 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '2' - name: SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.3 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '3' - name: PC2 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.4 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '4' - name: SWITCH1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.5 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '5' - name: SWITCH2 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.6 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '6' - name: SWITCH3 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.7 - software_state: GOOD - file_system_state: GOOD -- item_type: LINK - id: '7' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '4' -- item_type: LINK - id: '8' - name: link2 - bandwidth: 1000000000 - source: '4' - destination: '2' -- item_type: LINK - id: '9' - name: link3 - bandwidth: 1000000000 - source: '2' - destination: '5' -- item_type: LINK - id: '10' - name: link4 - bandwidth: 1000000000 - source: '2' - destination: '6' -- item_type: LINK - id: '11' - name: link5 - bandwidth: 1000000000 - source: '5' - destination: '3' -- item_type: LINK - id: '12' - name: link6 - bandwidth: 1000000000 - source: '6' - destination: '3' -- item_type: GREEN_IER - id: '13' - start_step: 1 - end_step: 128 - load: 100000 - protocol: TCP - port: '80' - source: '3' - destination: '2' - mission_criticality: 5 -- item_type: RED_POL - id: '14' - start_step: 50 - end_step: 50 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_IER - id: '15' - start_step: 60 - end_step: 100 - load: 1000000 - protocol: TCP - port: '80' - source: '1' - destination: '2' - mission_criticality: 0 -- item_type: RED_POL - id: '16' - start_step: 80 - end_step: 80 - targetNodeId: '2' - initiator: IER - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: ACL_RULE - id: '17' - permission: ALLOW - source: ANY - destination: ANY - protocol: ANY - port: ANY - position: 0 diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml deleted file mode 100644 index e91859d2..00000000 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml +++ /dev/null @@ -1,366 +0,0 @@ -- item_type: PORTS - ports_list: - - port: '80' -- item_type: SERVICES - service_list: - - name: TCP -- item_type: NODE - node_id: '1' - name: PC1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.11 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '2' - name: PC2 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.12 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '3' - name: PC3 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.13 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '4' - name: PC4 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.20.14 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '5' - name: SWITCH1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.2 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '6' - name: IDS - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.4 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '7' - name: SWITCH2 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.3 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '8' - name: LOP1 - node_class: SERVICE - node_type: LOP - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.12 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '9' - name: SERVER1 - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.14 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '10' - name: SERVER2 - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.20.15 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: LINK - id: '11' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '5' -- item_type: LINK - id: '12' - name: link2 - bandwidth: 1000000000 - source: '2' - destination: '5' -- item_type: LINK - id: '13' - name: link3 - bandwidth: 1000000000 - source: '3' - destination: '5' -- item_type: LINK - id: '14' - name: link4 - bandwidth: 1000000000 - source: '4' - destination: '5' -- item_type: LINK - id: '15' - name: link5 - bandwidth: 1000000000 - source: '5' - destination: '6' -- item_type: LINK - id: '16' - name: link6 - bandwidth: 1000000000 - source: '5' - destination: '8' -- item_type: LINK - id: '17' - name: link7 - bandwidth: 1000000000 - source: '6' - destination: '7' -- item_type: LINK - id: '18' - name: link8 - bandwidth: 1000000000 - source: '8' - destination: '7' -- item_type: LINK - id: '19' - name: link9 - bandwidth: 1000000000 - source: '7' - destination: '9' -- item_type: LINK - id: '20' - name: link10 - bandwidth: 1000000000 - source: '7' - destination: '10' -- item_type: GREEN_IER - id: '21' - start_step: 1 - end_step: 128 - load: 100000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - mission_criticality: 2 -- item_type: GREEN_IER - id: '22' - start_step: 1 - end_step: 128 - load: 100000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - mission_criticality: 2 -- item_type: GREEN_IER - id: '23' - start_step: 1 - end_step: 128 - load: 100000 - protocol: TCP - port: '80' - source: '9' - destination: '3' - mission_criticality: 5 -- item_type: GREEN_IER - id: '24' - start_step: 1 - end_step: 128 - load: 100000 - protocol: TCP - port: '80' - source: '4' - destination: '10' - mission_criticality: 2 -- item_type: ACL_RULE - id: '25' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.10.14 - protocol: TCP - port: 80 - position: 0 -- item_type: ACL_RULE - id: '26' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.10.14 - protocol: TCP - port: 80 - position: 1 -- item_type: ACL_RULE - id: '27' - permission: ALLOW - source: 192.168.10.13 - destination: 192.168.10.14 - protocol: TCP - port: 80 - position: 2 -- item_type: ACL_RULE - id: '28' - permission: ALLOW - source: 192.168.20.14 - destination: 192.168.20.15 - protocol: TCP - port: 80 - position: 3 -- item_type: ACL_RULE - id: '29' - permission: ALLOW - source: 192.168.10.14 - destination: 192.168.10.13 - protocol: TCP - port: 80 - position: 4 -- item_type: ACL_RULE - id: '30' - permission: DENY - source: 192.168.10.11 - destination: 192.168.20.15 - protocol: TCP - port: 80 - position: 5 -- item_type: ACL_RULE - id: '31' - permission: DENY - source: 192.168.10.12 - destination: 192.168.20.15 - protocol: TCP - port: 80 - position: 6 -- item_type: ACL_RULE - id: '32' - permission: DENY - source: 192.168.10.13 - destination: 192.168.20.15 - protocol: TCP - port: 80 - position: 7 -- item_type: ACL_RULE - id: '33' - permission: DENY - source: 192.168.20.14 - destination: 192.168.10.14 - protocol: TCP - port: 80 - position: 8 -- item_type: RED_POL - id: '34' - start_step: 20 - end_step: 20 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '35' - start_step: 20 - end_step: 20 - targetNodeId: '2' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_IER - id: '36' - start_step: 30 - end_step: 128 - load: 440000000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - mission_criticality: 0 -- item_type: RED_IER - id: '37' - start_step: 30 - end_step: 128 - load: 440000000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - mission_criticality: 0 -- item_type: RED_POL - id: '38' - start_step: 30 - end_step: 30 - targetNodeId: '9' - initiator: IER - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml deleted file mode 100644 index 453b6abb..00000000 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml +++ /dev/null @@ -1,164 +0,0 @@ -- item_type: PORTS - ports_list: - - port: '80' -- item_type: SERVICES - service_list: - - name: TCP -- item_type: NODE - node_id: '1' - name: PC1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.2 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '2' - name: PC2 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.3 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '3' - name: SWITCH1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.1 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '4' - name: SERVER1 - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.4 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: LINK - id: '5' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '3' -- item_type: LINK - id: '6' - name: link2 - bandwidth: 1000000000 - source: '2' - destination: '3' -- item_type: LINK - id: '7' - name: link3 - bandwidth: 1000000000 - source: '3' - destination: '4' -- item_type: GREEN_IER - id: '8' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '1' - destination: '4' - mission_criticality: 1 -- item_type: GREEN_IER - id: '9' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '2' - destination: '4' - mission_criticality: 1 -- item_type: GREEN_IER - id: '10' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '4' - destination: '2' - mission_criticality: 5 -- item_type: ACL_RULE - id: '11' - permission: ALLOW - source: 192.168.1.2 - destination: 192.168.1.4 - protocol: TCP - port: 80 - position: 0 -- item_type: ACL_RULE - id: '12' - permission: ALLOW - source: 192.168.1.3 - destination: 192.168.1.4 - protocol: TCP - port: 80 - position: 1 -- item_type: ACL_RULE - id: '13' - permission: ALLOW - source: 192.168.1.4 - destination: 192.168.1.3 - protocol: TCP - port: 80 - position: 2 -- item_type: RED_POL - id: '14' - start_step: 20 - end_step: 20 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_IER - id: '15' - start_step: 30 - end_step: 256 - load: 10000000 - protocol: TCP - port: '80' - source: '1' - destination: '4' - mission_criticality: 0 -- item_type: RED_POL - id: '16' - start_step: 40 - end_step: 40 - targetNodeId: '4' - initiator: IER - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml deleted file mode 100644 index 96596514..00000000 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml +++ /dev/null @@ -1,546 +0,0 @@ -- item_type: PORTS - ports_list: - - port: '80' - - port: '1433' - - port: '53' -- item_type: SERVICES - service_list: - - name: TCP - - name: TCP_SQL - - name: UDP -- item_type: NODE - node_id: '1' - name: CLIENT_1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.11 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- item_type: NODE - node_id: '2' - name: CLIENT_2 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.10.12 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: NODE - node_id: '3' - name: SWITCH_1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.10.1 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '4' - name: SECURITY_SUITE - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.10 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- item_type: NODE - node_id: '5' - name: MANAGEMENT_CONSOLE - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.12 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- item_type: NODE - node_id: '6' - name: SWITCH_2 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.2.1 - software_state: GOOD - file_system_state: GOOD -- item_type: NODE - node_id: '7' - name: WEB_SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.2.10 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: TCP_SQL - port: '1433' - state: GOOD -- item_type: NODE - node_id: '8' - name: DATABASE_SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.2.14 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: TCP_SQL - port: '1433' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- item_type: NODE - node_id: '9' - name: BACKUP_SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.2.16 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- item_type: LINK - id: '10' - name: LINK_1 - bandwidth: 1000000000 - source: '1' - destination: '3' -- item_type: LINK - id: '11' - name: LINK_2 - bandwidth: 1000000000 - source: '2' - destination: '3' -- item_type: LINK - id: '12' - name: LINK_3 - bandwidth: 1000000000 - source: '3' - destination: '4' -- item_type: LINK - id: '13' - name: LINK_4 - bandwidth: 1000000000 - source: '3' - destination: '5' -- item_type: LINK - id: '14' - name: LINK_5 - bandwidth: 1000000000 - source: '4' - destination: '6' -- item_type: LINK - id: '15' - name: LINK_6 - bandwidth: 1000000000 - source: '5' - destination: '6' -- item_type: LINK - id: '16' - name: LINK_7 - bandwidth: 1000000000 - source: '6' - destination: '7' -- item_type: LINK - id: '17' - name: LINK_8 - bandwidth: 1000000000 - source: '6' - destination: '8' -- item_type: LINK - id: '18' - name: LINK_9 - bandwidth: 1000000000 - source: '6' - destination: '9' -- item_type: GREEN_IER - id: '19' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '1' - destination: '7' - mission_criticality: 5 -- item_type: GREEN_IER - id: '20' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '7' - destination: '1' - mission_criticality: 5 -- item_type: GREEN_IER - id: '21' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '2' - destination: '7' - mission_criticality: 5 -- item_type: GREEN_IER - id: '22' - start_step: 1 - end_step: 256 - load: 10000 - protocol: TCP - port: '80' - source: '7' - destination: '2' - mission_criticality: 5 -- item_type: GREEN_IER - id: '23' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP_SQL - port: '1433' - source: '7' - destination: '8' - mission_criticality: 5 -- item_type: GREEN_IER - id: '24' - start_step: 1 - end_step: 256 - load: 100000 - protocol: TCP_SQL - port: '1433' - source: '8' - destination: '7' - mission_criticality: 5 -- item_type: GREEN_IER - id: '25' - start_step: 1 - end_step: 256 - load: 50000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - mission_criticality: 2 -- item_type: GREEN_IER - id: '26' - start_step: 1 - end_step: 256 - load: 50000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - mission_criticality: 2 -- item_type: GREEN_IER - id: '27' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '7' - mission_criticality: 1 -- item_type: GREEN_IER - id: '28' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '7' - destination: '5' - mission_criticality: 1 -- item_type: GREEN_IER - id: '29' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '8' - mission_criticality: 1 -- item_type: GREEN_IER - id: '30' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '8' - destination: '5' - mission_criticality: 1 -- item_type: GREEN_IER - id: '31' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '9' - mission_criticality: 1 -- item_type: GREEN_IER - id: '32' - start_step: 1 - end_step: 256 - load: 5000 - protocol: TCP - port: '80' - source: '9' - destination: '5' - mission_criticality: 1 -- item_type: ACL_RULE - id: '33' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.10 - protocol: ANY - port: ANY - position: 0 -- item_type: ACL_RULE - id: '34' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.14 - protocol: ANY - port: ANY - position: 1 -- item_type: ACL_RULE - id: '35' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.14 - protocol: ANY - port: ANY - position: 2 -- item_type: ACL_RULE - id: '36' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.10 - protocol: ANY - port: ANY - position: 3 -- item_type: ACL_RULE - id: '37' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.10.11 - protocol: ANY - port: ANY - position: 4 -- item_type: ACL_RULE - id: '38' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.10.12 - protocol: ANY - port: ANY - position: 5 -- item_type: ACL_RULE - id: '39' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.2.14 - protocol: ANY - port: ANY - position: 6 -- item_type: ACL_RULE - id: '40' - permission: ALLOW - source: 192.168.2.14 - destination: 192.168.2.10 - protocol: ANY - port: ANY - position: 7 -- item_type: ACL_RULE - id: '41' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.16 - protocol: ANY - port: ANY - position: 8 -- item_type: ACL_RULE - id: '42' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.16 - protocol: ANY - port: ANY - position: 9 -- item_type: ACL_RULE - id: '43' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.10 - protocol: ANY - port: ANY - position: 10 -- item_type: ACL_RULE - id: '44' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.14 - protocol: ANY - port: ANY - position: 11 -- item_type: ACL_RULE - id: '45' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.16 - protocol: ANY - port: ANY - position: 12 -- item_type: ACL_RULE - id: '46' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.1.12 - protocol: ANY - port: ANY - position: 13 -- item_type: ACL_RULE - id: '47' - permission: ALLOW - source: 192.168.2.14 - destination: 192.168.1.12 - protocol: ANY - port: ANY - position: 14 -- item_type: ACL_RULE - id: '48' - permission: ALLOW - source: 192.168.2.16 - destination: 192.168.1.12 - protocol: ANY - port: ANY - position: 15 -- item_type: ACL_RULE - id: '49' - permission: DENY - source: ANY - destination: ANY - protocol: ANY - port: ANY - position: 16 -- item_type: RED_POL - id: '50' - start_step: 50 - end_step: 50 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: UDP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_IER - id: '51' - start_step: 75 - end_step: 105 - load: 10000 - protocol: UDP - port: '53' - source: '1' - destination: '8' - mission_criticality: 0 -- item_type: RED_POL - id: '52' - start_step: 100 - end_step: 100 - targetNodeId: '8' - initiator: IER - type: SERVICE - protocol: UDP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '53' - start_step: 105 - end_step: 105 - targetNodeId: '8' - initiator: SERVICE - type: FILE - protocol: NA - state: CORRUPT - sourceNodeId: '8' - sourceNodeService: UDP - sourceNodeServiceState: COMPROMISED -- item_type: RED_POL - id: '54' - start_step: 105 - end_step: 105 - targetNodeId: '8' - initiator: SERVICE - type: SERVICE - protocol: TCP_SQL - state: COMPROMISED - sourceNodeId: '8' - sourceNodeService: UDP - sourceNodeServiceState: COMPROMISED -- item_type: RED_POL - id: '55' - start_step: 125 - end_step: 125 - targetNodeId: '7' - initiator: SERVICE - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: '8' - sourceNodeService: TCP_SQL - sourceNodeServiceState: COMPROMISED diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml deleted file mode 100644 index db4ed692..00000000 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ /dev/null @@ -1,168 +0,0 @@ -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: null - -# Set whether the agent evaluation will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# observation space -observation_space: - flatten: true - components: - - name: NODE_LINK_TABLE - - name: NODE_STATUSES - - name: LINK_TRAFFIC_LEVELS - - name: ACCESS_CONTROL_LIST - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 1 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 10 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# Implicit ACL firewall rule at end of ACL list to be the default action (ALLOW or DENY) -implicit_acl_rule: DENY -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 30 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -0.001 -off_should_be_resetting: -0.0005 -on_should_be_off: -0.0002 -on_should_be_resetting: -0.0005 -resetting_should_be_on: -0.0005 -resetting_should_be_off: -0.0002 -resetting: -0.0003 -# Node Software or Service State -good_should_be_patching: 0.0002 -good_should_be_compromised: 0.0005 -good_should_be_overwhelmed: 0.0005 -patching_should_be_good: -0.0005 -patching_should_be_compromised: 0.0002 -patching_should_be_overwhelmed: 0.0002 -patching: -0.0003 -compromised_should_be_good: -0.002 -compromised_should_be_patching: -0.002 -compromised_should_be_overwhelmed: -0.002 -compromised: -0.002 -overwhelmed_should_be_good: -0.002 -overwhelmed_should_be_patching: -0.002 -overwhelmed_should_be_compromised: -0.002 -overwhelmed: -0.002 -# Node File System State -good_should_be_repairing: 0.0002 -good_should_be_restoring: 0.0002 -good_should_be_corrupt: 0.0005 -good_should_be_destroyed: 0.001 -repairing_should_be_good: -0.0005 -repairing_should_be_restoring: 0.0002 -repairing_should_be_corrupt: 0.0002 -repairing_should_be_destroyed: 0.0000 -repairing: -0.0003 -restoring_should_be_good: -0.001 -restoring_should_be_repairing: -0.0002 -restoring_should_be_corrupt: 0.0001 -restoring_should_be_destroyed: 0.0002 -restoring: -0.0006 -corrupt_should_be_good: -0.001 -corrupt_should_be_repairing: -0.001 -corrupt_should_be_restoring: -0.001 -corrupt_should_be_destroyed: 0.0002 -corrupt: -0.001 -destroyed_should_be_good: -0.002 -destroyed_should_be_repairing: -0.002 -destroyed_should_be_restoring: -0.002 -destroyed_should_be_corrupt: -0.002 -destroyed: -0.002 -scanning: -0.0002 -# IER status -red_ier_running: -0.0005 -green_ier_blocked: -0.001 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/src/primaite/config/lay_down_config.py b/src/primaite/config/lay_down_config.py deleted file mode 100644 index fe3e3429..00000000 --- a/src/primaite/config/lay_down_config.py +++ /dev/null @@ -1,141 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from logging import Logger -from pathlib import Path -from typing import Any, Dict, Final, List, Union - -import yaml - -from primaite import getLogger, PRIMAITE_PATHS - -_LOGGER: Logger = getLogger(__name__) - -_EXAMPLE_LAY_DOWN: Final[Path] = PRIMAITE_PATHS.user_config_path / "example_config" / "lay_down" - - -def convert_legacy_lay_down_config(legacy_config: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Convert a legacy lay down config to the new format. - - :param legacy_config: A legacy lay down config. - """ - field_conversion_map = { - "itemType": "item_type", - "portsList": "ports_list", - "serviceList": "service_list", - "baseType": "node_class", - "nodeType": "node_type", - "hardwareState": "hardware_state", - "softwareState": "software_state", - "startStep": "start_step", - "endStep": "end_step", - "fileSystemState": "file_system_state", - "ipAddress": "ip_address", - "missionCriticality": "mission_criticality", - } - new_config = [] - for item in legacy_config: - if "itemType" in item: - if item["itemType"] in ["ACTIONS", "STEPS"]: - continue - new_dict = {} - for key in item.keys(): - conversion_key = field_conversion_map.get(key) - if key == "id" and "itemType" in item: - if item["itemType"] == "NODE": - conversion_key = "node_id" - if conversion_key: - new_dict[conversion_key] = item[key] - else: - new_dict[key] = item[key] - new_config.append(new_dict) - return new_config - - -def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict: - """ - Read in a lay down config yaml file. - - :param file_path: The config file path. - :param legacy_file: True if the config file is legacy format, otherwise False. - :return: The lay down config as a dict. - :raises ValueError: If the file_path does not exist. - """ - if not isinstance(file_path, Path): - file_path = Path(file_path) - if file_path.exists(): - with open(file_path, "r") as file: - config = yaml.safe_load(file) - _LOGGER.debug(f"Loading lay down config file: {file_path}") - if legacy_file: - try: - config = convert_legacy_lay_down_config(config) - except KeyError: - msg = ( - f"Failed to convert lay down config file {file_path} " - f"from legacy format. Attempting to use file as is." - ) - _LOGGER.error(msg) - return config - msg = f"Cannot load the lay down config as it does not exist: {file_path}" - _LOGGER.error(msg) - raise ValueError(msg) - - -def ddos_basic_one_config_path() -> Path: - """ - The path to the example lay_down_config_1_DDOS_basic.yaml file. - - :return: The file path. - """ - path = _EXAMPLE_LAY_DOWN / "lay_down_config_1_DDOS_basic.yaml" - if not path.exists(): - msg = "Example config not found. Please run 'primaite setup'" - _LOGGER.critical(msg) - raise FileNotFoundError(msg) - - return path - - -def ddos_basic_two_config_path() -> Path: - """ - The path to the example lay_down_config_2_DDOS_basic.yaml file. - - :return: The file path. - """ - path = _EXAMPLE_LAY_DOWN / "lay_down_config_2_DDOS_basic.yaml" - if not path.exists(): - msg = "Example config not found. Please run 'primaite setup'" - _LOGGER.critical(msg) - raise FileNotFoundError(msg) - - return path - - -def dos_very_basic_config_path() -> Path: - """ - The path to the example lay_down_config_3_DOS_very_basic.yaml file. - - :return: The file path. - """ - path = _EXAMPLE_LAY_DOWN / "lay_down_config_3_DOS_very_basic.yaml" - if not path.exists(): - msg = "Example config not found. Please run 'primaite setup'" - _LOGGER.critical(msg) - raise FileNotFoundError(msg) - - return path - - -def data_manipulation_config_path() -> Path: - """ - The path to the example lay_down_config_5_data_manipulation.yaml file. - - :return: The file path. - """ - path = _EXAMPLE_LAY_DOWN / "lay_down_config_5_data_manipulation.yaml" - if not path.exists(): - msg = "Example config not found. Please run 'primaite setup'" - _LOGGER.critical(msg) - raise FileNotFoundError(msg) - - return path diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py deleted file mode 100644 index f81bb6f7..00000000 --- a/src/primaite/config/training_config.py +++ /dev/null @@ -1,438 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from __future__ import annotations - -from dataclasses import dataclass, field -from logging import Logger -from pathlib import Path -from typing import Any, Dict, Final, Optional, Union - -import yaml - -from primaite import getLogger, PRIMAITE_PATHS -from primaite.common.enums import ( - ActionType, - AgentFramework, - AgentIdentifier, - DeepLearningFramework, - HardCodedAgentView, - RulePermissionType, - SB3OutputVerboseLevel, - SessionType, -) - -_LOGGER: Logger = getLogger(__name__) - -_EXAMPLE_TRAINING: Final[Path] = PRIMAITE_PATHS.user_config_path / "example_config" / "training" - - -def main_training_config_path() -> Path: - """ - The path to the example training_config_main.yaml file. - - :return: The file path. - """ - path = _EXAMPLE_TRAINING / "training_config_main.yaml" - if not path.exists(): - msg = "Example config not found. Please run 'primaite setup'" - _LOGGER.critical(msg) - raise FileNotFoundError(msg) - - return path - - -@dataclass() -class TrainingConfig: - """The Training Config class.""" - - agent_framework: AgentFramework = AgentFramework.SB3 - "The AgentFramework" - - deep_learning_framework: DeepLearningFramework = DeepLearningFramework.TF - "The DeepLearningFramework" - - agent_identifier: AgentIdentifier = AgentIdentifier.PPO - "The AgentIdentifier" - - hard_coded_agent_view: HardCodedAgentView = HardCodedAgentView.FULL - "The view the deterministic hard-coded agent has of the environment" - - random_red_agent: bool = False - "Creates Random Red Agent Attacks" - - action_type: ActionType = ActionType.ANY - "The ActionType to use" - - num_train_episodes: int = 10 - "The number of episodes to train over during an training session" - - num_train_steps: int = 256 - "The number of steps in an episode during an training session" - - num_eval_episodes: int = 1 - "The number of episodes to train over during an evaluation session" - - num_eval_steps: int = 256 - "The number of steps in an episode during an evaluation session" - - checkpoint_every_n_episodes: int = 5 - "The agent will save a checkpoint every n episodes" - - observation_space: dict = field(default_factory=lambda: {"components": [{"name": "NODE_LINK_TABLE"}]}) - "The observation space config dict" - - time_delay: int = 10 - "The delay between steps (ms). Applies to generic agents only" - - # file - session_type: SessionType = SessionType.TRAIN - "The type of PrimAITE session to run" - - load_agent: bool = False - "Determine whether to load an agent from file" - - agent_load_file: Optional[str] = None - "File path and file name of agent if you're loading one in" - - # Environment - observation_space_high_value: int = 1000000000 - "The high value for the observation space" - - sb3_output_verbose_level: SB3OutputVerboseLevel = SB3OutputVerboseLevel.NONE - "Stable Baselines3 learn/eval output verbosity level" - - implicit_acl_rule: RulePermissionType = RulePermissionType.DENY - "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." - - max_number_acl_rules: int = 30 - "Sets a limit for number of acl rules allowed in the list and environment." - - # Reward values - # Generic - all_ok: float = 0 - - # Node Hardware State - off_should_be_on: float = -0.001 - off_should_be_resetting: float = -0.0005 - on_should_be_off: float = -0.0002 - on_should_be_resetting: float = -0.0005 - resetting_should_be_on: float = -0.0005 - resetting_should_be_off: float = -0.0002 - resetting: float = -0.0003 - - # Node Software or Service State - good_should_be_patching: float = 0.0002 - good_should_be_compromised: float = 0.0005 - good_should_be_overwhelmed: float = 0.0005 - patching_should_be_good: float = -0.0005 - patching_should_be_compromised: float = 0.0002 - patching_should_be_overwhelmed: float = 0.0002 - patching: float = -0.0003 - compromised_should_be_good: float = -0.002 - compromised_should_be_patching: float = -0.002 - compromised_should_be_overwhelmed: float = -0.002 - compromised: float = -0.002 - overwhelmed_should_be_good: float = -0.002 - overwhelmed_should_be_patching: float = -0.002 - overwhelmed_should_be_compromised: float = -0.002 - overwhelmed: float = -0.002 - - # Node File System State - good_should_be_repairing: float = 0.0002 - good_should_be_restoring: float = 0.0002 - good_should_be_corrupt: float = 0.0005 - good_should_be_destroyed: float = 0.001 - repairing_should_be_good: float = -0.0005 - repairing_should_be_restoring: float = 0.0002 - repairing_should_be_corrupt: float = 0.0002 - repairing_should_be_destroyed: float = 0.0000 - repairing: float = -0.0003 - restoring_should_be_good: float = -0.001 - restoring_should_be_repairing: float = -0.0002 - restoring_should_be_corrupt: float = 0.0001 - restoring_should_be_destroyed: float = 0.0002 - restoring: float = -0.0006 - corrupt_should_be_good: float = -0.001 - corrupt_should_be_repairing: float = -0.001 - corrupt_should_be_restoring: float = -0.001 - corrupt_should_be_destroyed: float = 0.0002 - corrupt: float = -0.001 - destroyed_should_be_good: float = -0.002 - destroyed_should_be_repairing: float = -0.002 - destroyed_should_be_restoring: float = -0.002 - destroyed_should_be_corrupt: float = -0.002 - destroyed: float = -0.002 - scanning: float = -0.0002 - - # IER status - red_ier_running: float = -0.0005 - green_ier_blocked: float = -0.001 - - # Patching / Reset durations - os_patching_duration: int = 5 - "The time taken to patch the OS" - - node_reset_duration: int = 5 - "The time taken to reset a node (hardware)" - - node_booting_duration: int = 3 - "The Time taken to turn on the node" - - node_shutdown_duration: int = 2 - "The time taken to turn off the node" - - service_patching_duration: int = 5 - "The time taken to patch a service" - - file_system_repairing_limit: int = 5 - "The time take to repair the file system" - - file_system_restoring_limit: int = 5 - "The time take to restore the file system" - - file_system_scanning_limit: int = 5 - "The time taken to scan the file system" - - deterministic: bool = False - "If true, the training will be deterministic" - - seed: Optional[int] = None - "The random number generator seed to be used while training the agent" - - @classmethod - def from_dict(cls, config_dict: Dict[str, Any]) -> TrainingConfig: - """ - Create an instance of TrainingConfig from a dict. - - :param config_dict: The training config dict. - :return: The instance of TrainingConfig. - """ - field_enum_map = { - "agent_framework": AgentFramework, - "deep_learning_framework": DeepLearningFramework, - "agent_identifier": AgentIdentifier, - "action_type": ActionType, - "session_type": SessionType, - "sb3_output_verbose_level": SB3OutputVerboseLevel, - "hard_coded_agent_view": HardCodedAgentView, - "implicit_acl_rule": RulePermissionType, - } - - # convert the string representation of enums into the actual enum values themselves? - for key, value in field_enum_map.items(): - if key in config_dict: - config_dict[key] = value[config_dict[key]] - - return TrainingConfig(**config_dict) - - def to_dict(self, json_serializable: bool = True) -> Dict: - """ - Serialise the ``TrainingConfig`` as dict. - - :param json_serializable: If True, Enums are converted to their - string name. - :return: The ``TrainingConfig`` as a dict. - """ - data = self.__dict__ - if json_serializable: - data["agent_framework"] = self.agent_framework.name - data["deep_learning_framework"] = self.deep_learning_framework.name - data["agent_identifier"] = self.agent_identifier.name - data["action_type"] = self.action_type.name - data["sb3_output_verbose_level"] = self.sb3_output_verbose_level.name - data["session_type"] = self.session_type.name - data["hard_coded_agent_view"] = self.hard_coded_agent_view.name - data["implicit_acl_rule"] = self.implicit_acl_rule.name - - return data - - def __str__(self) -> str: - obs_str = ",".join([c["name"] for c in self.observation_space["components"]]) - tc = f"{self.agent_framework}, " - # if self.agent_framework is AgentFramework.RLLIB: - # tc += f"{self.deep_learning_framework}, " - tc += f"{self.agent_identifier}, " - if self.agent_identifier is AgentIdentifier.HARDCODED: - tc += f"{self.hard_coded_agent_view}, " - tc += f"{self.action_type}, " - tc += f"observation_space={obs_str}, " - if self.session_type is SessionType.TRAIN: - tc += f"{self.num_train_episodes} episodes @ " - tc += f"{self.num_train_steps} steps" - elif self.session_type is SessionType.EVAL: - tc += f"{self.num_eval_episodes} episodes @ " - tc += f"{self.num_eval_steps} steps" - else: - tc += f"Training: {self.num_eval_episodes} episodes @ " - tc += f"{self.num_eval_steps} steps" - tc += f"Evaluation: {self.num_eval_episodes} episodes @ " - tc += f"{self.num_eval_steps} steps" - return tc - - -def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConfig: - """ - Read in a training config yaml file. - - :param file_path: The config file path. - :param legacy_file: True if the config file is legacy format, otherwise - False. - :return: An instance of - :class:`~primaite.config.training_config.TrainingConfig`. - :raises ValueError: If the file_path does not exist. - :raises TypeError: When the TrainingConfig object cannot be created - using the values from the config file read from ``file_path``. - """ - if not isinstance(file_path, Path): - file_path = Path(file_path) - if file_path.exists(): - with open(file_path, "r") as file: - config = yaml.safe_load(file) - _LOGGER.debug(f"Loading training config file: {file_path}") - if legacy_file: - try: - config = convert_legacy_training_config_dict(config) - - except KeyError as e: - msg = ( - f"Failed to convert training config file {file_path} " - f"from legacy format. Attempting to use file as is." - ) - _LOGGER.error(msg) - raise e - try: - return TrainingConfig.from_dict(config) - except TypeError as e: - msg = f"Error when creating an instance of {TrainingConfig} " f"from the training config file {file_path}" - _LOGGER.critical(msg, exc_info=True) - raise e - msg = f"Cannot load the training config as it does not exist: {file_path}" - _LOGGER.error(msg) - raise ValueError(msg) - - -def convert_legacy_training_config_dict( - legacy_config_dict: Dict[str, Any], - agent_framework: AgentFramework = AgentFramework.SB3, - agent_identifier: AgentIdentifier = AgentIdentifier.PPO, - action_type: ActionType = ActionType.ANY, - num_train_steps: int = 256, - num_eval_steps: int = 256, - num_train_episodes: int = 10, - num_eval_episodes: int = 1, -) -> Dict[str, Any]: - """ - Convert a legacy training config dict to the new format. - - :param legacy_config_dict: A legacy training config dict. - :param agent_framework: The agent framework to use as legacy training - configs don't have agent_framework values. - :param agent_identifier: The red agent identifier to use as legacy - training configs don't have agent_identifier values. - :param action_type: The action space type to set as legacy training configs - don't have action_type values. - :param num_train_steps: The number of train steps to set as legacy training configs - don't have num_train_steps values. - :param num_eval_steps: The number of eval steps to set as legacy training configs - don't have num_eval_steps values. - :param num_train_episodes: The number of train episodes to set as legacy training configs - don't have num_train_episodes values. - :param num_eval_episodes: The number of eval episodes to set as legacy training configs - don't have num_eval_episodes values. - :return: The converted training config dict. - """ - config_dict = { - "agent_framework": agent_framework.name, - "agent_identifier": agent_identifier.name, - "action_type": action_type.name, - "num_train_steps": num_train_steps, - "num_eval_steps": num_eval_steps, - "num_train_episodes": num_train_episodes, - "num_eval_episodes": num_eval_episodes, - "sb3_output_verbose_level": SB3OutputVerboseLevel.INFO.name, - } - session_type_map = {"TRAINING": "TRAIN", "EVALUATION": "EVAL"} - legacy_config_dict["sessionType"] = session_type_map[legacy_config_dict["sessionType"]] - for legacy_key, value in legacy_config_dict.items(): - new_key = _get_new_key_from_legacy(legacy_key) - if new_key: - config_dict[new_key] = value - return config_dict - - -def _get_new_key_from_legacy(legacy_key: str) -> Optional[str]: - """ - Maps legacy training config keys to the new format keys. - - :param legacy_key: A legacy training config key. - :return: The mapped key. - """ - key_mapping = { - "agentIdentifier": None, - "numEpisodes": "num_train_episodes", - "numSteps": "num_train_steps", - "timeDelay": "time_delay", - "configFilename": None, - "sessionType": "session_type", - "loadAgent": "load_agent", - "agentLoadFile": "agent_load_file", - "observationSpaceHighValue": "observation_space_high_value", - "allOk": "all_ok", - "offShouldBeOn": "off_should_be_on", - "offShouldBeResetting": "off_should_be_resetting", - "onShouldBeOff": "on_should_be_off", - "onShouldBeResetting": "on_should_be_resetting", - "resettingShouldBeOn": "resetting_should_be_on", - "resettingShouldBeOff": "resetting_should_be_off", - "resetting": "resetting", - "goodShouldBePatching": "good_should_be_patching", - "goodShouldBeCompromised": "good_should_be_compromised", - "goodShouldBeOverwhelmed": "good_should_be_overwhelmed", - "patchingShouldBeGood": "patching_should_be_good", - "patchingShouldBeCompromised": "patching_should_be_compromised", - "patchingShouldBeOverwhelmed": "patching_should_be_overwhelmed", - "patching": "patching", - "compromisedShouldBeGood": "compromised_should_be_good", - "compromisedShouldBePatching": "compromised_should_be_patching", - "compromisedShouldBeOverwhelmed": "compromised_should_be_overwhelmed", - "compromised": "compromised", - "overwhelmedShouldBeGood": "overwhelmed_should_be_good", - "overwhelmedShouldBePatching": "overwhelmed_should_be_patching", - "overwhelmedShouldBeCompromised": "overwhelmed_should_be_compromised", - "overwhelmed": "overwhelmed", - "goodShouldBeRepairing": "good_should_be_repairing", - "goodShouldBeRestoring": "good_should_be_restoring", - "goodShouldBeCorrupt": "good_should_be_corrupt", - "goodShouldBeDestroyed": "good_should_be_destroyed", - "repairingShouldBeGood": "repairing_should_be_good", - "repairingShouldBeRestoring": "repairing_should_be_restoring", - "repairingShouldBeCorrupt": "repairing_should_be_corrupt", - "repairingShouldBeDestroyed": "repairing_should_be_destroyed", - "repairing": "repairing", - "restoringShouldBeGood": "restoring_should_be_good", - "restoringShouldBeRepairing": "restoring_should_be_repairing", - "restoringShouldBeCorrupt": "restoring_should_be_corrupt", - "restoringShouldBeDestroyed": "restoring_should_be_destroyed", - "restoring": "restoring", - "corruptShouldBeGood": "corrupt_should_be_good", - "corruptShouldBeRepairing": "corrupt_should_be_repairing", - "corruptShouldBeRestoring": "corrupt_should_be_restoring", - "corruptShouldBeDestroyed": "corrupt_should_be_destroyed", - "corrupt": "corrupt", - "destroyedShouldBeGood": "destroyed_should_be_good", - "destroyedShouldBeRepairing": "destroyed_should_be_repairing", - "destroyedShouldBeRestoring": "destroyed_should_be_restoring", - "destroyedShouldBeCorrupt": "destroyed_should_be_corrupt", - "destroyed": "destroyed", - "scanning": "scanning", - "redIerRunning": "red_ier_running", - "greenIerBlocked": "green_ier_blocked", - "osPatchingDuration": "os_patching_duration", - "nodeResetDuration": "node_reset_duration", - "nodeBootingDuration": "node_booting_duration", - "nodeShutdownDuration": "node_shutdown_duration", - "servicePatchingDuration": "service_patching_duration", - "fileSystemRepairingLimit": "file_system_repairing_limit", - "fileSystemRestoringLimit": "file_system_restoring_limit", - "fileSystemScanningLimit": "file_system_scanning_limit", - } - return key_mapping[legacy_key] diff --git a/src/primaite/data_viz/__init__.py b/src/primaite/data_viz/__init__.py deleted file mode 100644 index 260579da..00000000 --- a/src/primaite/data_viz/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Utility to generate plots of sessions metrics after PrimAITE.""" -from enum import Enum - - -class PlotlyTemplate(Enum): - """The built-in plotly templates.""" - - PLOTLY = "plotly" - PLOTLY_WHITE = "plotly_white" - PLOTLY_DARK = "plotly_dark" - GGPLOT2 = "ggplot2" - SEABORN = "seaborn" - SIMPLE_WHITE = "simple_white" - NONE = "none" diff --git a/src/primaite/data_viz/session_plots.py b/src/primaite/data_viz/session_plots.py deleted file mode 100644 index 37750353..00000000 --- a/src/primaite/data_viz/session_plots.py +++ /dev/null @@ -1,73 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from pathlib import Path -from typing import Dict, Optional, Union - -import plotly.graph_objects as go -import polars as pl -import yaml -from plotly.graph_objs import Figure - -from primaite import PRIMAITE_PATHS - - -def get_plotly_config() -> Dict: - """Get the plotly config from primaite_config.yaml.""" - with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: - primaite_config = yaml.safe_load(file) - return primaite_config["session"]["outputs"]["plots"] - - -def plot_av_reward_per_episode( - av_reward_per_episode_csv: Union[str, Path], - title: Optional[str] = None, - subtitle: Optional[str] = None, -) -> Figure: - """ - Plot the average reward per episode from a csv session output. - - :param av_reward_per_episode_csv: The average reward per episode csv - file path. - :param title: The plot title. This is optional. - :param subtitle: The plot subtitle. This is optional. - :return: The plot as an instance of ``plotly.graph_objs._figure.Figure``. - """ - df = pl.read_csv(av_reward_per_episode_csv) - - if title: - if subtitle: - title = f"{title}
{subtitle}" - else: - if subtitle: - title = subtitle - - config = get_plotly_config() - layout = go.Layout( - autosize=config["size"]["auto_size"], - width=config["size"]["width"], - height=config["size"]["height"], - ) - # Create the line graph with a colored line - fig = go.Figure(layout=layout) - fig.update_layout(template=config["template"]) - fig.add_trace( - go.Scatter( - x=df["Episode"], - y=df["Average Reward"], - mode="lines", - name="Mean Reward per Episode", - ) - ) - - # Set the layout of the graph - fig.update_layout( - xaxis={ - "title": "Episode", - "type": "linear", - "rangeslider": {"visible": config["range_slider"]}, - }, - yaxis={"title": "Average Reward"}, - title=title, - showlegend=False, - ) - - return fig diff --git a/src/primaite/environment/__init__.py b/src/primaite/environment/__init__.py deleted file mode 100644 index f0fd21b9..00000000 --- a/src/primaite/environment/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Gym/Gymnasium environment for RL agents consisting of a simulated computer network.""" diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py deleted file mode 100644 index 73b9e998..00000000 --- a/src/primaite/environment/observations.py +++ /dev/null @@ -1,735 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Module for handling configurable observation spaces in PrimAITE.""" -import logging -from abc import ABC, abstractmethod -from logging import Logger -from typing import Dict, Final, List, Tuple, TYPE_CHECKING, Union - -import numpy as np -from gymnasium import spaces - -from primaite.acl.acl_rule import ACLRule -from primaite.common.enums import FileSystemState, HardwareState, RulePermissionType, SoftwareState -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.service_node import ServiceNode - -# This dependency is only needed for type hints, -# TYPE_CHECKING is False at runtime and True when typecheckers are performing typechecking -# Therefore, this avoids circular dependency problem. -if TYPE_CHECKING: - from primaite.environment.primaite_env import Primaite - - -_LOGGER: Logger = logging.getLogger(__name__) - - -class AbstractObservationComponent(ABC): - """Represents a part of the PrimAITE observation space.""" - - @abstractmethod - def __init__(self, env: "Primaite") -> None: - """ - Initialise observation component. - - :param env: Primaite training environment. - :type env: Primaite - """ - _LOGGER.info(f"Initialising {self} observation component") - self.env: "Primaite" = env - self.space: spaces.Space - self.current_observation: np.ndarray # type might be too restrictive? - self.structure: List[str] - return NotImplemented - - @abstractmethod - def update(self) -> None: - """Update the observation based on the current state of the environment.""" - self.current_observation = NotImplemented - - @abstractmethod - def generate_structure(self) -> List[str]: - """Return a list of labels for the components of the flattened observation space.""" - return NotImplemented - - -class NodeLinkTable(AbstractObservationComponent): - """ - Table with nodes and links as rows and hardware/software status as cols. - - This will create the observation space formatted as a table of integers. - There is one row per node, followed by one row per link. - The number of columns is 4 plus one per service. They are: - - * node/link ID - * node hardware status / 0 for links - * node operating system status (if active/service) / 0 for links - * node file system status (active/service only) / 0 for links - * node service1 status / traffic load from that service for links - * node service2 status / traffic load from that service for links - * ... - * node serviceN status / traffic load from that service for links - - For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be - ``(12, 7)`` - """ - - _FIXED_PARAMETERS: int = 4 - _MAX_VAL: int = 1_000_000_000 - _DATA_TYPE: type = np.int64 - - def __init__(self, env: "Primaite") -> None: - """ - Initialise a NodeLinkTable observation space component. - - :param env: Training environment. - :type env: Primaite - """ - super().__init__(env) - - # 1. Define the shape of your observation space component - num_items = self.env.num_links + self.env.num_nodes - num_columns = self.env.num_services + self._FIXED_PARAMETERS - observation_shape = (num_items, num_columns) - - # 2. Create Observation space - self.space = spaces.Box( - low=0, - high=self._MAX_VAL, - shape=observation_shape, - dtype=self._DATA_TYPE, - ) - - # 3. Initialise Observation with zeroes - self.current_observation = np.zeros(observation_shape, dtype=self._DATA_TYPE) - - self.structure = self.generate_structure() - - def update(self) -> None: - """ - Update the observation based on current environment state. - - The structure of the observation space is described in :class:`.NodeLinkTable` - """ - item_index = 0 - nodes = self.env.nodes - links = self.env.links - # Do nodes first - for _, node in nodes.items(): - self.current_observation[item_index][0] = int(node.node_id) - self.current_observation[item_index][1] = node.hardware_state.value - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - self.current_observation[item_index][2] = node.software_state.value - self.current_observation[item_index][3] = node.file_system_state_observed.value - else: - self.current_observation[item_index][2] = 0 - self.current_observation[item_index][3] = 0 - service_index = 4 - if isinstance(node, ServiceNode): - for service in self.env.services_list: - if node.has_service(service): - self.current_observation[item_index][service_index] = node.get_service_state(service).value - else: - self.current_observation[item_index][service_index] = 0 - service_index += 1 - else: - # Not a service node - for service in self.env.services_list: - self.current_observation[item_index][service_index] = 0 - service_index += 1 - item_index += 1 - - # Now do links - for _, link in links.items(): - self.current_observation[item_index][0] = int(link.get_id()) - self.current_observation[item_index][1] = 0 - self.current_observation[item_index][2] = 0 - self.current_observation[item_index][3] = 0 - protocol_list = link.get_protocol_list() - protocol_index = 0 - for protocol in protocol_list: - self.current_observation[item_index][protocol_index + 4] = protocol.get_load() - protocol_index += 1 - item_index += 1 - - def generate_structure(self) -> List[str]: - """Return a list of labels for the components of the flattened observation space.""" - nodes = self.env.nodes.values() - links = self.env.links.values() - - structure = [] - - for i, node in enumerate(nodes): - node_id = node.node_id - node_labels = [ - f"node_{node_id}_id", - f"node_{node_id}_hardware_status", - f"node_{node_id}_os_status", - f"node_{node_id}_fs_status", - ] - for j, serv in enumerate(self.env.services_list): - node_labels.append(f"node_{node_id}_service_{serv}_status") - - structure.extend(node_labels) - - for i, link in enumerate(links): - link_id = link.id - link_labels = [ - f"link_{link_id}_id", - f"link_{link_id}_n/a", - f"link_{link_id}_n/a", - f"link_{link_id}_n/a", - ] - for j, serv in enumerate(self.env.services_list): - link_labels.append(f"link_{link_id}_service_{serv}_load") - - structure.extend(link_labels) - return structure - - -class NodeStatuses(AbstractObservationComponent): - """ - Flat list of nodes' hardware, OS, file system, and service states. - - The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by - integers. - Each node has 3 elements plus 1 per service. It will have the following structure: - .. code-block:: - - [ - node1 hardware state, - node1 OS state, - node1 file system state, - node1 service1 state, - node1 service2 state, - node1 serviceN state (one for each service), - node2 hardware state, - node2 OS state, - node2 file system state, - node2 service1 state, - node2 service2 state, - node2 serviceN state (one for each service), - ... - ] - """ - - _DATA_TYPE: type = np.int64 - - def __init__(self, env: "Primaite") -> None: - """ - Initialise a NodeStatuses observation component. - - :param env: Training environment. - :type env: Primaite - """ - super().__init__(env) - - # 1. Define the shape of your observation space component - node_shape = [ - len(HardwareState) + 1, - len(SoftwareState) + 1, - len(FileSystemState) + 1, - ] - services_shape = [len(SoftwareState) + 1] * self.env.num_services - node_shape = node_shape + services_shape - - shape = node_shape * self.env.num_nodes - # 2. Create Observation space - self.space = spaces.MultiDiscrete(shape) - - # 3. Initialise observation with zeroes - self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) - self.structure = self.generate_structure() - - def update(self) -> None: - """ - Update the observation based on current environment state. - - The structure of the observation space is described in :class:`.NodeStatuses` - """ - obs = [] - for _, node in self.env.nodes.items(): - hardware_state = node.hardware_state.value - software_state = 0 - file_system_state = 0 - service_states = [0] * self.env.num_services - - if isinstance(node, ActiveNode): - software_state = node.software_state.value - file_system_state = node.file_system_state_observed.value - - if isinstance(node, ServiceNode): - for i, service in enumerate(self.env.services_list): - if node.has_service(service): - service_states[i] = node.get_service_state(service).value - obs.extend( - [ - hardware_state, - software_state, - file_system_state, - *service_states, - ] - ) - self.current_observation[:] = obs - - def generate_structure(self) -> List[str]: - """Return a list of labels for the components of the flattened observation space.""" - services = self.env.services_list - - structure = [] - - for _, node in self.env.nodes.items(): - node_id = node.node_id - structure.append(f"node_{node_id}_hardware_state_NONE") - for state in HardwareState: - structure.append(f"node_{node_id}_hardware_state_{state.name}") - structure.append(f"node_{node_id}_software_state_NONE") - for state in SoftwareState: - structure.append(f"node_{node_id}_software_state_{state.name}") - structure.append(f"node_{node_id}_file_system_state_NONE") - for state in FileSystemState: - structure.append(f"node_{node_id}_file_system_state_{state.name}") - for service in services: - structure.append(f"node_{node_id}_service_{service}_state_NONE") - for state in SoftwareState: - structure.append(f"node_{node_id}_service_{service}_state_{state.name}") - return structure - - -class LinkTrafficLevels(AbstractObservationComponent): - """ - Flat list of traffic levels encoded into banded categories. - - For each link, total traffic or traffic per service is encoded into a categorical value. - For example, if ``quantisation_levels=5``, the traffic levels represent these values: - - * 0 = No traffic (0% of bandwidth) - * 1 = No traffic (0%-33% of bandwidth) - * 2 = No traffic (33%-66% of bandwidth) - * 3 = No traffic (66%-100% of bandwidth) - * 4 = No traffic (100% of bandwidth) - - .. note:: - The lowest category always corresponds to no traffic and the highest category to the link being at max capacity. - Any amount of traffic between 0% and 100% (exclusive) is divided evenly into the remaining categories. - - """ - - _DATA_TYPE: type = np.int64 - - def __init__( - self, - env: "Primaite", - combine_service_traffic: bool = False, - quantisation_levels: int = 5, - ) -> None: - """ - Initialise a LinkTrafficLevels observation component. - - :param env: The environment that forms the basis of the observations - :type env: Primaite - :param combine_service_traffic: Whether to consider total traffic on the link, or each protocol individually, - defaults to False - :type combine_service_traffic: bool, optional - :param quantisation_levels: How many bands to consider when converting the traffic amount to a categorical - value, defaults to 5 - :type quantisation_levels: int, optional - """ - if quantisation_levels < 3: - _msg = ( - f"quantisation_levels must be 3 or more because the lowest and highest levels are " - f"reserved for 0% and 100% link utilisation, got {quantisation_levels} instead. " - f"Resetting to default value (5)" - ) - _LOGGER.warning(_msg) - quantisation_levels = 5 - - super().__init__(env) - - self._combine_service_traffic: bool = combine_service_traffic - self._quantisation_levels: int = quantisation_levels - self._entries_per_link: int = 1 - - if not self._combine_service_traffic: - self._entries_per_link = self.env.num_services - - # 1. Define the shape of your observation space component - shape = [self._quantisation_levels] * self.env.num_links * self._entries_per_link - - # 2. Create Observation space - self.space = spaces.MultiDiscrete(shape) - - # 3. Initialise observation with zeroes - self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) - - self.structure = self.generate_structure() - - def update(self) -> None: - """ - Update the observation based on current environment state. - - The structure of the observation space is described in :class:`.LinkTrafficLevels` - """ - obs = [] - for _, link in self.env.links.items(): - bandwidth = link.bandwidth - if self._combine_service_traffic: - loads = [link.get_current_load()] - else: - loads = [protocol.get_load() for protocol in link.protocol_list] - - for load in loads: - if load <= 0: - traffic_level = 0 - elif load >= bandwidth: - traffic_level = self._quantisation_levels - 1 - else: - traffic_level = (load / bandwidth) // (1 / (self._quantisation_levels - 2)) + 1 - - obs.append(int(traffic_level)) - - self.current_observation[:] = obs - - def generate_structure(self) -> List[str]: - """Return a list of labels for the components of the flattened observation space.""" - structure = [] - for _, link in self.env.links.items(): - link_id = link.id - if self._combine_service_traffic: - protocols = ["overall"] - else: - protocols = [protocol.name for protocol in link.protocol_list] - - for p in protocols: - for i in range(self._quantisation_levels): - structure.append(f"link_{link_id}_{p}_traffic_level_{i}") - return structure - - -class AccessControlList(AbstractObservationComponent): - """Flat list of all the Access Control Rules in the Access Control List. - - The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by - integers. - - Each ACL Rule has 6 elements. It will have the following structure: - .. code-block:: - [ - acl_rule1 permission, - acl_rule1 source_ip, - acl_rule1 dest_ip, - acl_rule1 protocol, - acl_rule1 port, - acl_rule1 position, - acl_rule2 permission, - acl_rule2 source_ip, - acl_rule2 dest_ip, - acl_rule2 protocol, - acl_rule2 port, - acl_rule2 position, - ... - ] - - - Terms (for ACL Observation Space): - [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) - [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to Node IDs) - [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to Node IDs) - [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) - [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) - [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) - - NOTE: NA is Non-Applicable - this means the ACL Rule in the list is a NoneType and NOT an ACLRule object. - """ - - _DATA_TYPE: type = np.int64 - - def __init__(self, env: "Primaite") -> None: - """ - Initialise an AccessControlList observation component. - - :param env: The environment that forms the basis of the observations - :type env: Primaite - """ - super().__init__(env) - - # 1. Define the shape of your observation space component - # The NA and ANY types means that there are 2 extra items for Nodes, Services and Ports. - # Number of ACL rules incremented by 1 for positions starting at index 0. - acl_shape = [ - len(RulePermissionType), - len(env.nodes) + 2, - len(env.nodes) + 2, - len(env.services_list) + 2, - len(env.ports_list) + 2, - env.max_number_acl_rules, - ] - shape = acl_shape * self.env.max_number_acl_rules - - # 2. Create Observation space - self.space = spaces.MultiDiscrete(shape) - - # 3. Initialise observation with zeroes - self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) - - self.structure = self.generate_structure() - - def update(self) -> None: - """Update the observation based on current environment state. - - The structure of the observation space is described in :class:`.AccessControlList` - """ - obs = [] - - for index in range(0, len(self.env.acl.acl)): - acl_rule = self.env.acl.acl[index] - if isinstance(acl_rule, ACLRule): - permission = acl_rule.permission - source_ip = acl_rule.source_ip - dest_ip = acl_rule.dest_ip - protocol = acl_rule.protocol - port = acl_rule.port - position = index - # Map each ACL attribute from what it was to an integer to fit the observation space - source_ip_int = None - dest_ip_int = None - if permission == RulePermissionType.DENY: - permission_int = 1 - else: - permission_int = 2 - if source_ip == "ANY": - source_ip_int = 1 - else: - # Map Node ID (+ 1) to source IP address - nodes = list(self.env.nodes.values()) - for node in nodes: - if ( - isinstance(node, ServiceNode) or isinstance(node, ActiveNode) - ) and node.ip_address == source_ip: - source_ip_int = int(node.node_id) + 1 - break - if dest_ip == "ANY": - dest_ip_int = 1 - else: - # Map Node ID (+ 1) to dest IP address - # Index of Nodes start at 1 so + 1 is needed so NA can be added. - nodes = list(self.env.nodes.values()) - for node in nodes: - if ( - isinstance(node, ServiceNode) or isinstance(node, ActiveNode) - ) and node.ip_address == dest_ip: - dest_ip_int = int(node.node_id) + 1 - if protocol == "ANY": - protocol_int = 1 - else: - # Index of protocols and ports start from 0 so + 2 is needed to add NA and ANY - try: - protocol_int = self.env.services_list.index(protocol) + 2 - except AttributeError: - _LOGGER.info(f"Service {protocol} could not be found") - protocol_int = None - if port == "ANY": - port_int = 1 - else: - if port in self.env.ports_list: - port_int = self.env.ports_list.index(port) + 2 - else: - _LOGGER.info(f"Port {port} could not be found.") - port_int = None - # Add to current obs - obs.extend( - [ - permission_int, - source_ip_int, - dest_ip_int, - protocol_int, - port_int, - position, - ] - ) - - else: - # The Nothing or NA representation of 'NONE' ACL rules - obs.extend([0, 0, 0, 0, 0, 0]) - - self.current_observation[:] = obs - - def generate_structure(self) -> List[str]: - """Return a list of labels for the components of the flattened observation space.""" - structure = [] - for acl_rule in self.env.acl.acl: - acl_rule_id = self.env.acl.acl.index(acl_rule) - - for permission in RulePermissionType: - structure.append(f"acl_rule_{acl_rule_id}_permission_{permission.name}") - - structure.append(f"acl_rule_{acl_rule_id}_source_ip_ANY") - for node in self.env.nodes.keys(): - structure.append(f"acl_rule_{acl_rule_id}_source_ip_{node}") - - structure.append(f"acl_rule_{acl_rule_id}_dest_ip_ANY") - for node in self.env.nodes.keys(): - structure.append(f"acl_rule_{acl_rule_id}_dest_ip_{node}") - - structure.append(f"acl_rule_{acl_rule_id}_service_ANY") - for service in self.env.services_list: - structure.append(f"acl_rule_{acl_rule_id}_service_{service}") - - structure.append(f"acl_rule_{acl_rule_id}_port_ANY") - for port in self.env.ports_list: - structure.append(f"acl_rule_{acl_rule_id}_port_{port}") - - return structure - - -class ObservationsHandler: - """ - Component-based observation space handler. - - This allows users to configure observation spaces by mixing and matching components. Each component can also define - further parameters to make them more flexible. - """ - - _REGISTRY: Final[Dict[str, type]] = { - "NODE_LINK_TABLE": NodeLinkTable, - "NODE_STATUSES": NodeStatuses, - "LINK_TRAFFIC_LEVELS": LinkTrafficLevels, - "ACCESS_CONTROL_LIST": AccessControlList, - } - - def __init__(self) -> None: - """Initialise the observation handler.""" - self.registered_obs_components: List[AbstractObservationComponent] = [] - - # internal the observation space (unflattened version of space if flatten=True) - self._space: spaces.Space - # flattened version of the observation space - self._flat_space: spaces.Space - - self._observation: Union[Tuple[np.ndarray], np.ndarray] - # used for transactions and when flatten=true - self._flat_observation: np.ndarray - - def update_obs(self) -> None: - """Fetch fresh information about the environment.""" - current_obs = [] - for obs in self.registered_obs_components: - obs.update() - current_obs.append(obs.current_observation) - - if len(current_obs) == 1: - self._observation = current_obs[0] - else: - self._observation = tuple(current_obs) - self._flat_observation = spaces.flatten(self._space, self._observation) - - def register(self, obs_component: AbstractObservationComponent) -> None: - """ - Add a component for this handler to track. - - :param obs_component: The component to add. - :type obs_component: AbstractObservationComponent - """ - self.registered_obs_components.append(obs_component) - self.update_space() - - def deregister(self, obs_component: AbstractObservationComponent) -> None: - """ - Remove a component from this handler. - - :param obs_component: Which component to remove. It must exist within this object's - ``registered_obs_components`` attribute. - :type obs_component: AbstractObservationComponent - """ - self.registered_obs_components.remove(obs_component) - self.update_space() - - def update_space(self) -> None: - """Rebuild the handler's composite observation space from its components.""" - component_spaces = [] - for obs_comp in self.registered_obs_components: - component_spaces.append(obs_comp.space) - - # if there are multiple components, build a composite tuple space - if len(component_spaces) == 1: - self._space = component_spaces[0] - else: - self._space = spaces.Tuple(component_spaces) - if len(component_spaces) > 0: - self._flat_space = spaces.flatten_space(self._space) - else: - self._flat_space = spaces.Box(0, 1, (0,)) - - @property - def space(self) -> spaces.Space: - """Observation space, return the flattened version if flatten is True.""" - if len(self.registered_obs_components) > 1: - return self._flat_space - else: - return self._space - - @property - def current_observation(self) -> Union[np.ndarray, Tuple[np.ndarray]]: - """Current observation, return the flattened version if flatten is True.""" - if len(self.registered_obs_components) > 1: - return self._flat_observation - else: - return self._observation - - @classmethod - def from_config(cls, env: "Primaite", obs_space_config: dict) -> "ObservationsHandler": - """ - Parse a config dictinary, return a new observation handler populated with new observation component objects. - - The expected format for the config dictionary is: - - .. code-block:: python - - config = { - components: [ - { - "name": "" - }, - { - "name": "" - "options": {"opt1": val1, "opt2": val2} - }, - { - ... - }, - ] - } - - :return: Observation handler - :rtype: primaite.environment.observations.ObservationsHandler - """ - # Instantiate the handler - handler = cls() - - for component_cfg in obs_space_config["components"]: - # Figure out which class can instantiate the desired component - comp_type = component_cfg["name"] - comp_class = cls._REGISTRY[comp_type] - - # Create the component with options from the YAML - options = component_cfg.get("options") or {} - component = comp_class(env, **options) - - handler.register(component) - - handler.update_obs() - return handler - - def describe_structure(self) -> List[str]: - """ - Create a list of names for the features of the obs space. - - The order of labels follows the flattened version of the space. - """ - # as it turns out it's not possible to take the gym flattening function and apply it to our labels so we have - # to fake it. each component has to just hard-code the expected label order after flattening... - - labels = [] - for obs_comp in self.registered_obs_components: - labels.extend(obs_comp.structure) - - return labels diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py deleted file mode 100644 index a809772f..00000000 --- a/src/primaite/environment/primaite_env.py +++ /dev/null @@ -1,1408 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Main environment module containing the PRIMmary AI Training Evironment (Primaite) class.""" -import copy -import logging -import uuid as uuid -from logging import Logger -from pathlib import Path -from random import choice, randint, sample, uniform -from typing import Any, Dict, Final, List, Tuple, Union - -import networkx as nx -import numpy as np -from gym import Env, spaces -from matplotlib import pyplot as plt - -from primaite import getLogger -from primaite.acl.access_control_list import AccessControlList -from primaite.agents.utils import is_valid_acl_action_extra, is_valid_node_action -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import ( # AgentFramework, - ActionType, - AgentIdentifier, - FileSystemState, - HardwareState, - NodePOLInitiator, - NodePOLType, - NodeType, - ObservationType, - Priority, - SessionType, - SoftwareState, -) -from primaite.common.service import Service -from primaite.config import training_config -from primaite.config.lay_down_config import load -from primaite.config.training_config import TrainingConfig -from primaite.environment.observations import ObservationsHandler -from primaite.environment.reward import calculate_reward_function -from primaite.links.link import Link -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.node import Node -from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen -from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed -from primaite.nodes.passive_node import PassiveNode -from primaite.nodes.service_node import ServiceNode -from primaite.pol.green_pol import apply_iers, apply_node_pol -from primaite.pol.ier import IER -from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol -from primaite.transactions.transaction import Transaction -from primaite.utils.session_output_writer import SessionOutputWriter - -_LOGGER: Logger = getLogger(__name__) - - -class Primaite(Env): - """PRIMmary AI Training Evironment (Primaite) class.""" - - # Action Space contants - ACTION_SPACE_NODE_PROPERTY_VALUES: int = 5 - ACTION_SPACE_NODE_ACTION_VALUES: int = 4 - ACTION_SPACE_ACL_ACTION_VALUES: int = 3 - ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2 - - def __init__( - self, - training_config_path: Union[str, Path], - lay_down_config_path: Union[str, Path], - session_path: Path, - timestamp_str: str, - legacy_training_config: bool = False, - legacy_lay_down_config: bool = False, - ) -> None: - """ - The Primaite constructor. - - :param training_config_path: The training config filepath. - :param lay_down_config_path: The lay down config filepath. - :param session_path: The directory path the session is writing to. - :param timestamp_str: The session timestamp in the format: _. - :param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, - otherwise False. - """ - self.session_path: Final[Path] = session_path - self.timestamp_str: Final[str] = timestamp_str - self._training_config_path: Union[str, Path] = training_config_path - self._lay_down_config_path: Union[str, Path] = lay_down_config_path - self.legacy_training_config = legacy_training_config - self.legacy_lay_down_config = legacy_lay_down_config - - self.training_config: TrainingConfig = training_config.load(training_config_path, self.legacy_training_config) - _LOGGER.info(f"Using: {str(self.training_config)}") - - # Number of steps in an episode - self.episode_steps: int - if self.training_config.session_type == SessionType.TRAIN: - self.episode_steps = self.training_config.num_train_steps - elif self.training_config.session_type == SessionType.EVAL: - self.episode_steps = self.training_config.num_eval_steps - else: - self.episode_steps = self.training_config.num_train_steps - - super(Primaite, self).__init__() - - # The agent in use - self.agent_identifier: AgentIdentifier = self.training_config.agent_identifier - - # Create a dictionary to hold all the nodes - self.nodes: Dict[str, NodeUnion] = {} - - # Create a dictionary to hold a reference set of nodes - self.nodes_reference: Dict[str, NodeUnion] = {} - - # Create a dictionary to hold all the links - self.links: Dict[str, Link] = {} - - # Create a dictionary to hold a reference set of links - self.links_reference: Dict[str, Link] = {} - - # Create a dictionary to hold all the green IERs (this will come from an external source) - self.green_iers: Dict[str, IER] = {} - self.green_iers_reference: Dict[str, IER] = {} - - # Create a dictionary to hold all the node PoLs (this will come from an external source) - self.node_pol: Dict[str, NodeStateInstructionGreen] = {} - - # Create a dictionary to hold all the red agent IERs (this will come from an external source) - self.red_iers: Dict[str, IER] = {} - - # Create a dictionary to hold all the red agent node PoLs (this will come from an external source) - self.red_node_pol: Dict[str, NodeStateInstructionRed] = {} - - # Create the Access Control List - self.acl: AccessControlList = AccessControlList( - self.training_config.implicit_acl_rule, - self.training_config.max_number_acl_rules, - ) - # Sets limit for number of ACL rules in environment - self.max_number_acl_rules: int = self.training_config.max_number_acl_rules - - # Create a list of services (enums) - self.services_list: List[str] = [] - - # Create a list of ports - self.ports_list: List[str] = [] - - # Create graph (network) - self.network: nx.Graph = nx.MultiGraph() - - # Create a graph (network) reference - self.network_reference: nx.Graph = nx.MultiGraph() - - # Create step count - self.step_count: int = 0 - - self.total_step_count: int = 0 - """The total number of time steps completed.""" - - # Create step info dictionary - self.step_info: Dict[Any] = {} - - # Total reward - self.total_reward: float = 0 - - # Average reward - self.average_reward: float = 0 - - # Episode count - self.episode_count: int = 0 - - # Number of nodes - gets a value by examining the nodes dictionary after it's been populated - self.num_nodes: int = 0 - - # Number of links - gets a value by examining the links dictionary after it's been populated - self.num_links: int = 0 - - # Number of services - gets a value when config is loaded - self.num_services: int = 0 - - # Number of ports - gets a value when config is loaded - self.num_ports: int = 0 - - # The action type - # TODO: confirm type - self.action_type: int = 0 - - # TODO fix up with TrainingConfig - # stores the observation config from the yaml, default is NODE_LINK_TABLE - self.obs_config: dict = {"components": [{"name": "NODE_LINK_TABLE"}]} - if self.training_config.observation_space is not None: - self.obs_config = self.training_config.observation_space - - # Observation Handler manages the user-configurable observation space. - # It will be initialised later. - self.obs_handler: ObservationsHandler - - self._obs_space_description: List[str] = None - "The env observation space description for transactions writing" - - self.lay_down_config = load(self._lay_down_config_path, self.legacy_lay_down_config) - self.load_lay_down_config() - - # Store the node objects as node attributes - # (This is so we can access them as objects) - for node in self.network: - self.network.nodes[node]["self"] = node - - for node in self.network_reference: - self.network_reference.nodes[node]["self"] = node - - self.num_nodes = len(self.nodes) - self.num_links = len(self.links) - - # Visualise in PNG - try: - plt.tight_layout() - nx.draw_networkx(self.network, with_labels=True) - - file_path = session_path / f"network_{timestamp_str}.png" - plt.savefig(file_path, format="PNG") - plt.clf() - except Exception: - _LOGGER.error("Could not save network diagram", exc_info=True) - - # Initiate observation space - self.observation_space: spaces.Space - self.env_obs: np.ndarray - self.observation_space, self.env_obs = self.init_observations() - - # Define Action Space - depends on action space type (Node or ACL) - self.action_dict: Dict[int, List[int]] - self.action_space: spaces.Space - if self.training_config.action_type == ActionType.NODE: - _LOGGER.debug("Action space type NODE selected") - # Terms (for node action space): - # [0, num nodes] - node ID (0 = nothing, node ID) - # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, # noqa - # service state, file system state) # noqa - # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa - # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa - self.action_dict = self.create_node_action_dict() - self.action_space = spaces.Discrete(len(self.action_dict)) - elif self.training_config.action_type == ActionType.ACL: - _LOGGER.debug("Action space type ACL selected") - # Terms (for ACL action space): - # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - # [0, 1] - Permission (0 = DENY, 1 = ALLOW) - # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) - self.action_dict = self.create_acl_action_dict() - self.action_space = spaces.Discrete(len(self.action_dict)) - elif self.training_config.action_type == ActionType.ANY: - _LOGGER.debug("Action space type ANY selected - Node + ACL") - self.action_dict = self.create_node_and_acl_action_dict() - self.action_space = spaces.Discrete(len(self.action_dict)) - else: - _LOGGER.error(f"Invalid action type selected: {self.training_config.action_type}") - - self.episode_av_reward_writer: SessionOutputWriter = SessionOutputWriter( - self, transaction_writer=False, learning_session=True - ) - self.transaction_writer: SessionOutputWriter = SessionOutputWriter( - self, transaction_writer=True, learning_session=True - ) - - self.is_eval = False - - @property - def actual_episode_count(self) -> int: - """Shifts the episode_count by -1 for RLlib learning session.""" - # if self.training_config.agent_framework is AgentFramework.RLLIB and not self.is_eval: - # return self.episode_count - 1 - return self.episode_count - - def set_as_eval(self) -> None: - """Set the writers to write to eval directories.""" - self.episode_av_reward_writer = SessionOutputWriter(self, transaction_writer=False, learning_session=False) - self.transaction_writer = SessionOutputWriter(self, transaction_writer=True, learning_session=False) - self.episode_count = 0 - self.step_count = 0 - self.total_step_count = 0 - self.episode_steps = self.training_config.num_eval_steps - self.is_eval = True - - def _write_av_reward_per_episode(self) -> None: - if self.actual_episode_count > 0: - csv_data = self.actual_episode_count, self.average_reward - self.episode_av_reward_writer.write(csv_data) - - def reset(self) -> np.ndarray: - """ - AI Gym Reset function. - - Returns: - Environment observation space (reset) - """ - self._write_av_reward_per_episode() - self.episode_count += 1 - - # Don't need to reset links, as they are cleared and recalculated every - # step - - # Clear the ACL - self.init_acl() - - # Reset the node statuses and recreate the ACL from config - # Does this for both live and reference nodes - self.reset_environment() - - # Create a random red agent to use for this episode - if self.training_config.random_red_agent: - self._create_random_red_agent() - - # Reset counters and totals - self.total_reward = 0.0 - self.step_count = 0 - self.average_reward = 0.0 - - # Update observations space and return - self.update_environent_obs() - - return self.env_obs - - def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict]: - """ - AI Gym Step function. - - Args: - action: Action space from agent - - Returns: - env_obs: Observation space - reward: Reward value for this step - done: Indicates episode is complete if True - step_info: Additional information relating to this step - """ - # TEMP - done = False - self.step_count += 1 - self.total_step_count += 1 - - # Need to clear traffic on all links first - for link_key, link_value in self.links.items(): - link_value.clear_traffic() - - for link in self.links_reference.values(): - link.clear_traffic() - - # Create a Transaction (metric) object for this step - transaction = Transaction(self.agent_identifier, self.actual_episode_count, self.step_count) - # Load the initial observation space into the transaction - transaction.obs_space = self.obs_handler._flat_observation - - # Set the transaction obs space description - transaction.obs_space_description = self._obs_space_description - - # Load the action space into the transaction - transaction.action_space = copy.deepcopy(action) - - # 1. Implement Blue Action - self.interpret_action_and_apply(action) - # Take snapshots of nodes and links - self.nodes_post_blue = copy.deepcopy(self.nodes) - self.links_post_blue = copy.deepcopy(self.links) - - # 2. Perform any time-based activities (e.g. a component moving from patching to good) - self.apply_time_based_updates() - - # 3. Apply PoL - apply_node_pol(self.nodes, self.node_pol, self.step_count) # Node PoL - apply_iers( - self.network, - self.nodes, - self.links, - self.green_iers, - self.acl, - self.step_count, - ) # Network PoL - # Take snapshots of nodes and links - self.nodes_post_pol = copy.deepcopy(self.nodes) - self.links_post_pol = copy.deepcopy(self.links) - # Reference - apply_node_pol(self.nodes_reference, self.node_pol, self.step_count) # Node PoL - apply_iers( - self.network_reference, - self.nodes_reference, - self.links_reference, - self.green_iers_reference, - self.acl, - self.step_count, - ) # Network PoL - - # 4. Implement Red Action - apply_red_agent_iers( - self.network, - self.nodes, - self.links, - self.red_iers, - self.acl, - self.step_count, - ) - apply_red_agent_node_pol(self.nodes, self.red_iers, self.red_node_pol, self.step_count) - # Take snapshots of nodes and links - self.nodes_post_red = copy.deepcopy(self.nodes) - self.links_post_red = copy.deepcopy(self.links) - - # 5. Calculate reward signal (for RL) - reward = calculate_reward_function( - self.nodes_post_pol, - self.nodes_post_red, - self.nodes_reference, - self.green_iers, - self.green_iers_reference, - self.red_iers, - self.step_count, - self.training_config, - ) - _LOGGER.debug(f"Episode: {self.actual_episode_count}, " f"Step {self.step_count}, " f"Reward: {reward}") - self.total_reward += reward - if self.step_count == self.episode_steps: - self.average_reward = self.total_reward / self.step_count - if self.training_config.session_type is SessionType.EVAL: - # For evaluation, need to trigger the done value = True when - # step count is reached in order to prevent neverending episode - done = True - _LOGGER.info(f"Episode: {self.actual_episode_count}, " f"Average Reward: {self.average_reward}") - # Load the reward into the transaction - transaction.reward = reward - - # 6. Output Verbose - # self.output_link_status() - - # 7. Update env_obs - self.update_environent_obs() - - # Write transaction to file - if self.actual_episode_count > 0: - self.transaction_writer.write(transaction) - - # Return - return self.env_obs, reward, done, self.step_info - - def close(self) -> None: - """Override parent close and close writers.""" - # Close files if last episode/step - # if self.can_finish: - super().close() - - self.transaction_writer.close() - self.episode_av_reward_writer.close() - - def init_acl(self) -> None: - """Initialise the Access Control List.""" - self.acl.remove_all_rules() - - def output_link_status(self) -> None: - """Output the link status of all links to the console.""" - for link_key, link_value in self.links.items(): - _LOGGER.debug("Link ID: " + link_value.get_id()) - for protocol in link_value.protocol_list: - print(" Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load())) - - def interpret_action_and_apply(self, _action: int) -> None: - """ - Applies agent actions to the nodes and Access Control List. - - Args: - _action: The action space from the agent - """ - # At the moment, actions are only affecting nodes - if self.training_config.action_type == ActionType.NODE: - self.apply_actions_to_nodes(_action) - elif self.training_config.action_type == ActionType.ACL: - self.apply_actions_to_acl(_action) - elif len(self.action_dict[_action]) == 7: # ACL actions in multidiscrete form have len 7 - self.apply_actions_to_acl(_action) - elif len(self.action_dict[_action]) == 4: # Node actions in multdiscrete (array) from have len 4 - self.apply_actions_to_nodes(_action) - else: - logging.error("Invalid action type found") - - def apply_actions_to_nodes(self, _action: int) -> None: - """ - Applies agent actions to the nodes. - - Args: - _action: The action space from the agent - """ - readable_action = self.action_dict[_action] - node_id = readable_action[0] - node_property = readable_action[1] - property_action = readable_action[2] - service_index = readable_action[3] - - # Check that the action is requesting a valid node - try: - node = self.nodes[str(node_id)] - except Exception: - return - - if node_property == 0: - # This is the do nothing action - return - elif node_property == 1: - # This is an action on the node Hardware State - if property_action == 0: - # Do nothing - return - elif property_action == 1: - # Turn on (only applicable if it's OFF, not if it's patching) - if node.hardware_state == HardwareState.OFF: - node.turn_on() - elif property_action == 2: - # Turn off - node.turn_off() - elif property_action == 3: - # Reset (only applicable if it's ON) - if node.hardware_state == HardwareState.ON: - node.reset() - else: - return - elif node_property == 2: - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - # This is an action on the node Software State - if property_action == 0: - # Do nothing - return - elif property_action == 1: - # Patch (valid action if it's good or compromised) - node.software_state = SoftwareState.PATCHING - else: - # Node is not of Active or Service Type - return - elif node_property == 3: - # This is an action on a node Service State - if isinstance(node, ServiceNode): - # This is an action on a node Service State - if property_action == 0: - # Do nothing - return - elif property_action == 1: - # Patch (valid action if it's good or compromised) - node.set_service_state(self.services_list[service_index], SoftwareState.PATCHING) - else: - # Node is not of Service Type - return - elif node_property == 4: - # This is an action on a node file system state - if isinstance(node, ActiveNode): - if property_action == 0: - # Do nothing - return - elif property_action == 1: - # Scan - node.start_file_system_scan() - elif property_action == 2: - # Repair - # You cannot repair a destroyed file system - it needs restoring - if node.file_system_state_actual != FileSystemState.DESTROYED: - node.set_file_system_state(FileSystemState.REPAIRING) - elif property_action == 3: - # Restore - node.set_file_system_state(FileSystemState.RESTORING) - else: - # Node is not of Active Type - return - else: - return - - def apply_actions_to_acl(self, _action: int) -> None: - """ - Applies agent actions to the Access Control List [TO DO]. - - Args: - _action: The action space from the agent - """ - # Convert discrete value back to multidiscrete - readable_action = self.action_dict[_action] - - action_decision = readable_action[0] - action_permission = readable_action[1] - action_source_ip = readable_action[2] - action_destination_ip = readable_action[3] - action_protocol = readable_action[4] - action_port = readable_action[5] - acl_rule_position = readable_action[6] - - if action_decision == 0: - # It's decided to do nothing - return - else: - # It's decided to create a new ACL rule or remove an existing rule - # Permission value - if action_permission == 0: - acl_rule_permission = "DENY" - else: - acl_rule_permission = "ALLOW" - # Source IP value - if action_source_ip == 0: - acl_rule_source = "ANY" - else: - node = list(self.nodes.values())[action_source_ip - 1] - if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): - acl_rule_source = node.ip_address - else: - return - # Destination IP value - if action_destination_ip == 0: - acl_rule_destination = "ANY" - else: - node = list(self.nodes.values())[action_destination_ip - 1] - if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): - acl_rule_destination = node.ip_address - else: - return - # Protocol value - if action_protocol == 0: - acl_rule_protocol = "ANY" - else: - acl_rule_protocol = self.services_list[action_protocol - 1] - # Port value - if action_port == 0: - acl_rule_port = "ANY" - else: - acl_rule_port = self.ports_list[action_port - 1] - - # Now add or remove - if action_decision == 1: - # Add the rule - self.acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_rule_position, - ) - elif action_decision == 2: - # Remove the rule - self.acl.remove_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - ) - else: - return - - def apply_time_based_updates(self) -> None: - """ - Updates anything that needs to count down and then change state. - - e.g. reset / patching status - """ - for node_key, node in self.nodes.items(): - if node.hardware_state == HardwareState.RESETTING: - node.update_resetting_status() - else: - pass - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - node.update_file_system_state() - if node.software_state == SoftwareState.PATCHING: - node.update_os_patching_status() - else: - pass - else: - pass - if isinstance(node, ServiceNode): - node.update_services_patching_status() - else: - pass - - for node_key, node in self.nodes_reference.items(): - if node.hardware_state == HardwareState.RESETTING: - node.update_resetting_status() - else: - pass - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - node.update_file_system_state() - if node.software_state == SoftwareState.PATCHING: - node.update_os_patching_status() - else: - pass - else: - pass - if isinstance(node, ServiceNode): - node.update_services_patching_status() - else: - pass - - def init_observations(self) -> Tuple[spaces.Space, np.ndarray]: - """ - Create the environment's observation handler. - - :return: The observation space, initial observation (zeroed out array with the correct shape) - :rtype: Tuple[spaces.Space, np.ndarray] - """ - self.obs_handler = ObservationsHandler.from_config(self, self.obs_config) - - if not self._obs_space_description: - self._obs_space_description = self.obs_handler.describe_structure() - - return self.obs_handler.space, self.obs_handler.current_observation - - def update_environent_obs(self) -> None: - """Updates the observation space based on the node and link status.""" - self.obs_handler.update_obs() - self.env_obs = self.obs_handler.current_observation - - def load_lay_down_config(self) -> None: - """Loads config data in order to build the environment configuration.""" - for item in self.lay_down_config: - if item["item_type"] == "NODE": - # Create a node - self.create_node(item) - elif item["item_type"] == "LINK": - # Create a link - self.create_link(item) - elif item["item_type"] == "GREEN_IER": - # Create a Green IER - self.create_green_ier(item) - elif item["item_type"] == "GREEN_POL": - # Create a Green PoL - self.create_green_pol(item) - elif item["item_type"] == "RED_IER": - # Create a Red IER - self.create_red_ier(item) - elif item["item_type"] == "RED_POL": - # Create a Red PoL - self.create_red_pol(item) - elif item["item_type"] == "ACL_RULE": - # Create an ACL rule - self.create_acl_rule(item) - elif item["item_type"] == "SERVICES": - # Create the list of services - self.create_services_list(item) - elif item["item_type"] == "PORTS": - # Create the list of ports - self.create_ports_list(item) - else: - item_type = item["item_type"] - _LOGGER.error(f"Invalid item_type: {item_type}") - pass - - _LOGGER.info("Environment configuration loaded") - print("Environment configuration loaded") - - def create_node(self, item: Dict) -> None: - """ - Creates a node from config data. - - Args: - item: A config data item - """ - # All nodes have these parameters - node_id = item["node_id"] - node_name = item["name"] - node_class = item["node_class"] - node_type = NodeType[item["node_type"]] - node_priority = Priority[item["priority"]] - node_hardware_state = HardwareState[item["hardware_state"]] - - if node_class == "PASSIVE": - node = PassiveNode( - node_id, - node_name, - node_type, - node_priority, - node_hardware_state, - self.training_config, - ) - elif node_class == "ACTIVE": - # Active nodes have IP address, Software State and file system state - node_ip_address = item["ip_address"] - node_software_state = SoftwareState[item["software_state"]] - node_file_system_state = FileSystemState[item["file_system_state"]] - node = ActiveNode( - node_id, - node_name, - node_type, - node_priority, - node_hardware_state, - node_ip_address, - node_software_state, - node_file_system_state, - self.training_config, - ) - elif node_class == "SERVICE": - # Service nodes have IP address, Software State, file system state and list of services - node_ip_address = item["ip_address"] - node_software_state = SoftwareState[item["software_state"]] - node_file_system_state = FileSystemState[item["file_system_state"]] - node = ServiceNode( - node_id, - node_name, - node_type, - node_priority, - node_hardware_state, - node_ip_address, - node_software_state, - node_file_system_state, - self.training_config, - ) - node_services = item["services"] - for service in node_services: - service_protocol = service["name"] - service_port = service["port"] - service_state = SoftwareState[service["state"]] - node.add_service(Service(service_protocol, service_port, service_state)) - else: - # Bad formatting - pass - - # Copy the node for the reference version - node_ref = copy.deepcopy(node) - - # Add node to node dictionary - self.nodes[node_id] = node - - # Add reference node to reference node dictionary - self.nodes_reference[node_id] = node_ref - - # Add node to network - self.network.add_nodes_from([node]) - - # Add node to network (reference) - self.network_reference.add_nodes_from([node_ref]) - - def create_link(self, item: Dict) -> None: - """ - Creates a link from config data. - - Args: - item: A config data item - """ - link_id = item["id"] - link_name = item["name"] - link_bandwidth = item["bandwidth"] - link_source = item["source"] - link_destination = item["destination"] - - source_node: Node = self.nodes[link_source] - dest_node: Node = self.nodes[link_destination] - - # Add link to network - self.network.add_edge(source_node, dest_node, id=link_name) - - # Add link to link dictionary - self.links[link_name] = Link( - link_id, - link_bandwidth, - source_node.name, - dest_node.name, - self.services_list, - ) - - # Reference - source_node_ref: Node = self.nodes_reference[link_source] - dest_node_ref: Node = self.nodes_reference[link_destination] - - # Add link to network (reference) - self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name) - - # Add link to link dictionary (reference) - self.links_reference[link_name] = Link( - link_id, - link_bandwidth, - source_node_ref.name, - dest_node_ref.name, - self.services_list, - ) - - def create_green_ier(self, item: Dict) -> None: - """ - Creates a green IER from config data. - - Args: - item: A config data item - """ - ier_id = item["id"] - ier_start_step = item["start_step"] - ier_end_step = item["end_step"] - ier_load = item["load"] - ier_protocol = item["protocol"] - ier_port = item["port"] - ier_source = item["source"] - ier_destination = item["destination"] - ier_mission_criticality = item["mission_criticality"] - - # Create IER and add to green IER dictionary - self.green_iers[ier_id] = IER( - ier_id, - ier_start_step, - ier_end_step, - ier_load, - ier_protocol, - ier_port, - ier_source, - ier_destination, - ier_mission_criticality, - ) - self.green_iers_reference[ier_id] = IER( - ier_id, - ier_start_step, - ier_end_step, - ier_load, - ier_protocol, - ier_port, - ier_source, - ier_destination, - ier_mission_criticality, - ) - - def create_red_ier(self, item: Dict) -> None: - """ - Creates a red IER from config data. - - Args: - item: A config data item - """ - ier_id = item["id"] - ier_start_step = item["start_step"] - ier_end_step = item["end_step"] - ier_load = item["load"] - ier_protocol = item["protocol"] - ier_port = item["port"] - ier_source = item["source"] - ier_destination = item["destination"] - ier_mission_criticality = item["mission_criticality"] - - # Create IER and add to red IER dictionary - self.red_iers[ier_id] = IER( - ier_id, - ier_start_step, - ier_end_step, - ier_load, - ier_protocol, - ier_port, - ier_source, - ier_destination, - ier_mission_criticality, - ) - - def create_green_pol(self, item: Dict) -> None: - """ - Creates a green PoL object from config data. - - Args: - item: A config data item - """ - pol_id = item["id"] - pol_start_step = item["start_step"] - pol_end_step = item["end_step"] - pol_node = item["nodeId"] - pol_type = NodePOLType[item["type"]] - - # State depends on whether this is Operating, Software, file system or Service PoL type - if pol_type == NodePOLType.OPERATING: - pol_state = HardwareState[item["state"]] - pol_protocol = "" - elif pol_type == NodePOLType.FILE: - pol_state = FileSystemState[item["state"]] - pol_protocol = "" - else: - pol_protocol = item["protocol"] - pol_state = SoftwareState[item["state"]] - - self.node_pol[pol_id] = NodeStateInstructionGreen( - pol_id, - pol_start_step, - pol_end_step, - pol_node, - pol_type, - pol_protocol, - pol_state, - ) - - def create_red_pol(self, item: Dict) -> None: - """ - Creates a red PoL object from config data. - - Args: - item: A config data item - """ - pol_id = item["id"] - pol_start_step = item["start_step"] - pol_end_step = item["end_step"] - pol_target_node_id = item["targetNodeId"] - pol_initiator = NodePOLInitiator[item["initiator"]] - pol_type = NodePOLType[item["type"]] - pol_protocol = item["protocol"] - - # State depends on whether this is Operating, Software, file system or Service PoL type - if pol_type == NodePOLType.OPERATING: - pol_state = HardwareState[item["state"]] - elif pol_type == NodePOLType.FILE: - pol_state = FileSystemState[item["state"]] - else: - pol_state = SoftwareState[item["state"]] - - pol_source_node_id = item["sourceNodeId"] - pol_source_node_service = item["sourceNodeService"] - pol_source_node_service_state = item["sourceNodeServiceState"] - - self.red_node_pol[pol_id] = NodeStateInstructionRed( - pol_id, - pol_start_step, - pol_end_step, - pol_target_node_id, - pol_initiator, - pol_type, - pol_protocol, - pol_state, - pol_source_node_id, - pol_source_node_service, - pol_source_node_service_state, - ) - - def create_acl_rule(self, item: Dict) -> None: - """ - Creates an ACL rule from config data. - - Args: - item: A config data item - """ - acl_rule_permission = item["permission"] - acl_rule_source = item["source"] - acl_rule_destination = item["destination"] - acl_rule_protocol = item["protocol"] - acl_rule_port = item["port"] - acl_rule_position = item.get("position") - - self.acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_rule_position, - ) - - # TODO: confirm typehint using runtime - def create_services_list(self, services: Dict) -> None: - """ - Creates a list of services (enum) from config data. - - Args: - item: A config data item representing the services - """ - service_list = services["service_list"] - - for service in service_list: - service_name = service["name"] - self.services_list.append(service_name) - - # Set the number of services - self.num_services = len(self.services_list) - - def create_ports_list(self, ports: Dict) -> None: - """ - Creates a list of ports from config data. - - Args: - item: A config data item representing the ports - """ - ports_list = ports["ports_list"] - - for port in ports_list: - port_value = port["port"] - self.ports_list.append(port_value) - - # Set the number of ports - self.num_ports = len(self.ports_list) - - # TODO: this is not used anymore, write a ticket to delete it - def get_observation_info(self, observation_info: Dict) -> None: - """ - Extracts observation_info. - - :param observation_info: Config item that defines which type of observation space to use - :type observation_info: str - """ - self.observation_type = ObservationType[observation_info["type"]] - - # TODO: this is not used anymore, write a ticket to delete it. - def get_action_info(self, action_info: Dict) -> None: - """ - Extracts action_info. - - Args: - item: A config data item representing action info - """ - self.action_type = ActionType[action_info["type"]] - - def save_obs_config(self, obs_config: dict) -> None: - """ - Cache the config for the observation space. - - This is necessary as the observation space can't be built while reading the config, - it must be done after all the nodes, links, and services have been initialised. - - :param obs_config: Parsed config relating to the observation space. The format is described in - :py:meth:`primaite.environment.observations.ObservationsHandler.from_config` - :type obs_config: dict - """ - self.obs_config = obs_config - - def reset_environment(self) -> None: - """ - Resets environment. - - Uses config data config data in order to build the environment configuration. - """ - for item in self.lay_down_config: - if item["item_type"] == "NODE": - # Reset a node's state (normal and reference) - self.reset_node(item) - elif item["item_type"] == "ACL_RULE": - # Create an ACL rule (these are cleared on reset, so just need to recreate them) - self.create_acl_rule(item) - else: - # Do nothing (bad formatting or not relevant to reset) - pass - - # Reset the IER status so they are not running initially - # Green IERs - for ier_key, ier_value in self.green_iers.items(): - ier_value.set_is_running(False) - # Red IERs - for ier_key, ier_value in self.red_iers.items(): - ier_value.set_is_running(False) - - def reset_node(self, item: Dict) -> None: - """ - Resets the statuses of a node. - - Args: - item: A config data item - """ - # All nodes have these parameters - node_id = item["node_id"] - node_class = item["node_class"] - node_hardware_state: HardwareState = HardwareState[item["hardware_state"]] - - node: NodeUnion = self.nodes[node_id] - node_ref = self.nodes_reference[node_id] - - # Reset the hardware state (common for all node types) - node.hardware_state = node_hardware_state - node_ref.hardware_state = node_hardware_state - - if node_class == "ACTIVE": - # Active nodes have Software State - node_software_state = SoftwareState[item["software_state"]] - node_file_system_state = FileSystemState[item["file_system_state"]] - node.software_state = node_software_state - node_ref.software_state = node_software_state - node.set_file_system_state(node_file_system_state) - node_ref.set_file_system_state(node_file_system_state) - elif node_class == "SERVICE": - # Service nodes have Software State and list of services - node_software_state = SoftwareState[item["software_state"]] - node_file_system_state = FileSystemState[item["file_system_state"]] - node.software_state = node_software_state - node_ref.software_state = node_software_state - node.set_file_system_state(node_file_system_state) - node_ref.set_file_system_state(node_file_system_state) - # Update service states - node_services = item["services"] - for service in node_services: - service_protocol = service["name"] - service_state = SoftwareState[service["state"]] - # Update node service state - node.set_service_state(service_protocol, service_state) - # Update reference node service state - node_ref.set_service_state(service_protocol, service_state) - else: - # Bad formatting - pass - - def create_node_action_dict(self) -> Dict[int, List[int]]: - """ - Creates a dictionary mapping each possible discrete action to more readable multidiscrete action. - - Note: Only actions that have the potential to change the state exist in the mapping (except for key 0) - - example return: - {0: [1, 0, 0, 0], - 1: [1, 1, 1, 0], - 2: [1, 1, 2, 0], - 3: [1, 1, 3, 0], - 4: [1, 2, 1, 0], - 5: [1, 3, 1, 0], - ... - } - """ - # Terms (for node action space): - # [0, num nodes] - node ID (0 = nothing, node ID) - # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa - # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa - # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa - # reserve 0 action to be a nothing action - actions = {0: [1, 0, 0, 0]} - action_key = 1 - for node in range(1, self.num_nodes + 1): - # 4 node properties (NONE, OPERATING, OS, SERVICE) - for node_property in range(4): - # Node Actions either: - # (NONE, ON, OFF, RESET) - operating state OR (NONE, PATCH) - OS/service state - # Use MAX to ensure we get them all - for node_action in range(4): - for service_state in range(self.num_services): - action = [node, node_property, node_action, service_state] - # check to see if it's a nothing action (has no effect) - if is_valid_node_action(action): - actions[action_key] = action - action_key += 1 - - return actions - - def create_acl_action_dict(self) -> Dict[int, List[int]]: - """Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.""" - # Terms (for ACL action space): - # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - # [0, 1] - Permission (0 = DENY, 1 = ALLOW) - # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) - # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) - # reserve 0 action to be a nothing action - actions = {0: [0, 0, 0, 0, 0, 0, 0]} - - action_key = 1 - # 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE - for action_decision in range(3): - # 2 possible action permissions 0 = DENY, 1 = CREATE - for action_permission in range(2): - # Number of nodes + 1 (for any) - for source_ip in range(self.num_nodes + 1): - for dest_ip in range(self.num_nodes + 1): - for protocol in range(self.num_services + 1): - for port in range(self.num_ports + 1): - for position in range(self.max_number_acl_rules - 1): - action = [ - action_decision, - action_permission, - source_ip, - dest_ip, - protocol, - port, - position, - ] - # Check to see if it is an action we want to include as possible - # i.e. not a nothing action - if is_valid_acl_action_extra(action): - actions[action_key] = action - action_key += 1 - - return actions - - def create_node_and_acl_action_dict(self) -> Dict[int, List[int]]: - """ - Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action. - - The dictionary contains actions of both Node and ACL action types. - """ - node_action_dict = self.create_node_action_dict() - acl_action_dict = self.create_acl_action_dict() - - # Change node keys to not overlap with acl keys - # Only 1 nothing action (key 0) is required, remove the other - new_node_action_dict = {k + len(acl_action_dict) - 1: v for k, v in node_action_dict.items() if k != 0} - - # Combine the Node dict and ACL dict - combined_action_dict = {**acl_action_dict, **new_node_action_dict} - return combined_action_dict - - def _create_random_red_agent(self) -> None: - """Decide on random red agent for the episode to be called in env.reset().""" - # Reset the current red iers and red node pol - self.red_iers = {} - self.red_node_pol = {} - - # Decide how many nodes become compromised - node_list = list(self.nodes.values()) - computers = [node for node in node_list if node.node_type == NodeType.COMPUTER] - max_num_nodes_compromised = len(computers) # only computers can become compromised - # random select between 1 and max_num_nodes_compromised - num_nodes_to_compromise = randint(1, max_num_nodes_compromised) - - # Decide which of the nodes to compromise - nodes_to_be_compromised = sample(computers, num_nodes_to_compromise) - - # choose a random compromise node to be source of attacks - source_node = choice(nodes_to_be_compromised) - - # For each of the nodes to be compromised decide which step they become compromised - max_step_compromised = self.episode_steps // 2 # always compromise in first half of episode - - # Bandwidth for all links - bandwidths = [i.get_bandwidth() for i in list(self.links.values())] - - if len(bandwidths) < 1: - msg = "Random red agent cannot be used on a network without any links" - _LOGGER.error(msg) - raise Exception(msg) - - servers = [node for node in node_list if node.node_type == NodeType.SERVER] - - for n, node in enumerate(nodes_to_be_compromised): - # 1: Use Node PoL to set node to compromised - - _id = str(uuid.uuid4()) - _start_step = randint(2, max_step_compromised + 1) # step compromised - pol_service_name = choice(list(node.services.keys())) - - source_node_service = choice(list(source_node.services.values())) - - red_pol = NodeStateInstructionRed( - _id=_id, - _start_step=_start_step, - _end_step=_start_step, # only run for 1 step - _target_node_id=node.node_id, - _pol_initiator="DIRECT", - _pol_type=NodePOLType["SERVICE"], - pol_protocol=pol_service_name, - _pol_state=SoftwareState.COMPROMISED, - _pol_source_node_id=source_node.node_id, - _pol_source_node_service=source_node_service.name, - _pol_source_node_service_state=source_node_service.software_state, - ) - - self.red_node_pol[_id] = red_pol - - # 2: Launch the attack from compromised node - set the IER - - ier_id = str(uuid.uuid4()) - # Launch the attack after node is compromised, and not right at the end of the episode - ier_start_step = randint(_start_step + 2, int(self.episode_steps * 0.8)) - ier_end_step = self.episode_steps - - # Randomise the load, as a percentage of a random link bandwith - ier_load = uniform(0.4, 0.8) * choice(bandwidths) - ier_protocol = pol_service_name # Same protocol as compromised node - ier_service = node.services[pol_service_name] - ier_port = ier_service.port - ier_mission_criticality = 0 # Red IER will never be important to green agent success - # We choose a node to attack based on the first that applies: - # a. Green IERs, select dest node of the red ier based on dest node of green IER - # b. Attack a random server that doesn't have a DENY acl rule in default config - # c. Attack a random server - possible_ier_destinations = [ - ier.get_dest_node_id() - for ier in list(self.green_iers.values()) - if ier.get_source_node_id() == node.node_id - ] - if len(possible_ier_destinations) < 1: - for server in servers: - if not self.acl.is_blocked( - node.ip_address, - server.ip_address, - ier_service, - ier_port, - ): - possible_ier_destinations.append(server.node_id) - if len(possible_ier_destinations) < 1: - # If still none found choose from all servers - possible_ier_destinations = [server.node_id for server in servers] - ier_dest = choice(possible_ier_destinations) - self.red_iers[ier_id] = IER( - ier_id, - ier_start_step, - ier_end_step, - ier_load, - ier_protocol, - ier_port, - node.node_id, - ier_dest, - ier_mission_criticality, - ) - - overwhelm_pol = red_pol - overwhelm_pol.id = str(uuid.uuid4()) - overwhelm_pol.end_step = self.episode_steps - - # 3: Make sure the targetted node can be set to overwhelmed - with node pol - # # TODO remove duplicate red pol for same targetted service - must take into account start step - - o_pol_id = str(uuid.uuid4()) - o_red_pol = NodeStateInstructionRed( - _id=o_pol_id, - _start_step=ier_start_step, - _end_step=self.episode_steps, - _target_node_id=ier_dest, - _pol_initiator="DIRECT", - _pol_type=NodePOLType["SERVICE"], - pol_protocol=ier_protocol, - _pol_state=SoftwareState.OVERWHELMED, - _pol_source_node_id=source_node.node_id, - _pol_source_node_service=source_node_service.name, - _pol_source_node_service_state=source_node_service.software_state, - ) - self.red_node_pol[o_pol_id] = o_red_pol diff --git a/src/primaite/environment/reward.py b/src/primaite/environment/reward.py deleted file mode 100644 index aa9dc97d..00000000 --- a/src/primaite/environment/reward.py +++ /dev/null @@ -1,386 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Implements reward function.""" -from logging import Logger -from typing import Dict, TYPE_CHECKING, Union - -from primaite import getLogger -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import FileSystemState, HardwareState, SoftwareState -from primaite.common.service import Service -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.service_node import ServiceNode - -if TYPE_CHECKING: - from primaite.config.training_config import TrainingConfig - from primaite.pol.ier import IER - -_LOGGER: Logger = getLogger(__name__) - - -def calculate_reward_function( - initial_nodes: Dict[str, NodeUnion], - final_nodes: Dict[str, NodeUnion], - reference_nodes: Dict[str, NodeUnion], - green_iers: Dict[str, "IER"], - green_iers_reference: Dict[str, "IER"], - red_iers: Dict[str, "IER"], - step_count: int, - config_values: "TrainingConfig", -) -> float: - """ - Compares the states of the initial and final nodes/links to get a reward. - - Args: - initial_nodes: The nodes before red and blue agents take effect - final_nodes: The nodes after red and blue agents take effect - reference_nodes: The nodes if there had been no red or blue effect - green_iers: The green IERs (should be running) - red_iers: Should be stopeed (ideally) by the blue agent - step_count: current step - config_values: Config values - """ - reward_value: float = 0.0 - - # For each node, compare hardware state, SoftwareState, service states - for node_key, final_node in final_nodes.items(): - initial_node = initial_nodes[node_key] - reference_node = reference_nodes[node_key] - - # Hardware State - reward_value += score_node_operating_state(final_node, initial_node, reference_node, config_values) - - # Software State - if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode): - reward_value += score_node_os_state(final_node, initial_node, reference_node, config_values) - - # Service State - if isinstance(final_node, ServiceNode): - reward_value += score_node_service_state(final_node, initial_node, reference_node, config_values) - - # File System State - if isinstance(final_node, ActiveNode): - reward_value += score_node_file_system(final_node, initial_node, reference_node, config_values) - - # Go through each red IER - penalise if it is running - for ier_key, ier_value in red_iers.items(): - start_step = ier_value.get_start_step() - stop_step = ier_value.get_end_step() - if step_count >= start_step and step_count <= stop_step: - if ier_value.get_is_running(): - reward_value += config_values.red_ier_running - - # Go through each green IER - penalise if it's not running (weighted) - # but only if it's supposed to be running (it's running in reference) - for ier_key, ier_value in green_iers.items(): - reference_ier = green_iers_reference[ier_key] - start_step = ier_value.get_start_step() - stop_step = ier_value.get_end_step() - if step_count >= start_step and step_count <= stop_step: - reference_blocked = not reference_ier.get_is_running() - live_blocked = not ier_value.get_is_running() - ier_reward = config_values.green_ier_blocked * ier_value.get_mission_criticality() - - if live_blocked and not reference_blocked: - reward_value += ier_reward - elif live_blocked and reference_blocked: - _LOGGER.debug( - ( - f"IER {ier_key} is blocked in the reference and live environments. " - f"Penalty of {ier_reward} was NOT applied." - ) - ) - elif not live_blocked and reference_blocked: - _LOGGER.debug( - ( - f"IER {ier_key} is blocked in the reference env but not in the live one. " - f"Penalty of {ier_reward} was NOT applied." - ) - ) - return reward_value - - -def score_node_operating_state( - final_node: NodeUnion, initial_node: NodeUnion, reference_node: NodeUnion, config_values: "TrainingConfig" -) -> float: - """ - Calculates score relating to the hardware state of a node. - - Args: - final_node: The node after red and blue agents take effect - initial_node: The node before red and blue agents take effect - reference_node: The node if there had been no red or blue effect - config_values: Config values - """ - score: float = 0.0 - final_node_operating_state = final_node.hardware_state - reference_node_operating_state = reference_node.hardware_state - - if final_node_operating_state == reference_node_operating_state: - # All is well - we're no different from the reference situation - score += config_values.all_ok - else: - # We're different from the reference situation - # Need to compare reference and final (current) state of node (i.e. at every step) - if reference_node_operating_state == HardwareState.ON: - if final_node_operating_state == HardwareState.OFF: - score += config_values.off_should_be_on - elif final_node_operating_state == HardwareState.RESETTING: - score += config_values.resetting_should_be_on - else: - pass - elif reference_node_operating_state == HardwareState.OFF: - if final_node_operating_state == HardwareState.ON: - score += config_values.on_should_be_off - elif final_node_operating_state == HardwareState.RESETTING: - score += config_values.resetting_should_be_off - else: - pass - elif reference_node_operating_state == HardwareState.RESETTING: - if final_node_operating_state == HardwareState.ON: - score += config_values.on_should_be_resetting - elif final_node_operating_state == HardwareState.OFF: - score += config_values.off_should_be_resetting - elif final_node_operating_state == HardwareState.RESETTING: - score += config_values.resetting - else: - pass - else: - pass - - return score - - -def score_node_os_state( - final_node: Union[ActiveNode, ServiceNode], - initial_node: Union[ActiveNode, ServiceNode], - reference_node: Union[ActiveNode, ServiceNode], - config_values: "TrainingConfig", -) -> float: - """ - Calculates score relating to the Software State of a node. - - Args: - final_node: The node after red and blue agents take effect - initial_node: The node before red and blue agents take effect - reference_node: The node if there had been no red or blue effect - config_values: Config values - """ - score: float = 0.0 - final_node_os_state = final_node.software_state - reference_node_os_state = reference_node.software_state - - if final_node_os_state == reference_node_os_state: - # All is well - we're no different from the reference situation - score += config_values.all_ok - else: - # We're different from the reference situation - # Need to compare reference and final (current) state of node (i.e. at every step) - if reference_node_os_state == SoftwareState.GOOD: - if final_node_os_state == SoftwareState.PATCHING: - score += config_values.patching_should_be_good - elif final_node_os_state == SoftwareState.COMPROMISED: - score += config_values.compromised_should_be_good - else: - pass - elif reference_node_os_state == SoftwareState.PATCHING: - if final_node_os_state == SoftwareState.GOOD: - score += config_values.good_should_be_patching - elif final_node_os_state == SoftwareState.COMPROMISED: - score += config_values.compromised_should_be_patching - elif final_node_os_state == SoftwareState.PATCHING: - score += config_values.patching - else: - pass - elif reference_node_os_state == SoftwareState.COMPROMISED: - if final_node_os_state == SoftwareState.GOOD: - score += config_values.good_should_be_compromised - elif final_node_os_state == SoftwareState.PATCHING: - score += config_values.patching_should_be_compromised - elif final_node_os_state == SoftwareState.COMPROMISED: - score += config_values.compromised - else: - pass - else: - pass - - return score - - -def score_node_service_state( - final_node: ServiceNode, initial_node: ServiceNode, reference_node: ServiceNode, config_values: "TrainingConfig" -) -> float: - """ - Calculates score relating to the service state(s) of a node. - - Args: - final_node: The node after red and blue agents take effect - initial_node: The node before red and blue agents take effect - reference_node: The node if there had been no red or blue effect - config_values: Config values - """ - score: float = 0.0 - final_node_services: Dict[str, Service] = final_node.services - reference_node_services: Dict[str, Service] = reference_node.services - - for service_key, final_service in final_node_services.items(): - reference_service = reference_node_services[service_key] - final_service = final_node_services[service_key] - - if final_service.software_state == reference_service.software_state: - # All is well - we're no different from the reference situation - score += config_values.all_ok - else: - # We're different from the reference situation - # Need to compare reference and final state of node (i.e. at every step) - if reference_service.software_state == SoftwareState.GOOD: - if final_service.software_state == SoftwareState.PATCHING: - score += config_values.patching_should_be_good - elif final_service.software_state == SoftwareState.COMPROMISED: - score += config_values.compromised_should_be_good - elif final_service.software_state == SoftwareState.OVERWHELMED: - score += config_values.overwhelmed_should_be_good - else: - pass - elif reference_service.software_state == SoftwareState.PATCHING: - if final_service.software_state == SoftwareState.GOOD: - score += config_values.good_should_be_patching - elif final_service.software_state == SoftwareState.COMPROMISED: - score += config_values.compromised_should_be_patching - elif final_service.software_state == SoftwareState.OVERWHELMED: - score += config_values.overwhelmed_should_be_patching - elif final_service.software_state == SoftwareState.PATCHING: - score += config_values.patching - else: - pass - elif reference_service.software_state == SoftwareState.COMPROMISED: - if final_service.software_state == SoftwareState.GOOD: - score += config_values.good_should_be_compromised - elif final_service.software_state == SoftwareState.PATCHING: - score += config_values.patching_should_be_compromised - elif final_service.software_state == SoftwareState.COMPROMISED: - score += config_values.compromised - elif final_service.software_state == SoftwareState.OVERWHELMED: - score += config_values.overwhelmed_should_be_compromised - else: - pass - elif reference_service.software_state == SoftwareState.OVERWHELMED: - if final_service.software_state == SoftwareState.GOOD: - score += config_values.good_should_be_overwhelmed - elif final_service.software_state == SoftwareState.PATCHING: - score += config_values.patching_should_be_overwhelmed - elif final_service.software_state == SoftwareState.COMPROMISED: - score += config_values.compromised_should_be_overwhelmed - elif final_service.software_state == SoftwareState.OVERWHELMED: - score += config_values.overwhelmed - else: - pass - else: - pass - - return score - - -def score_node_file_system( - final_node: Union[ActiveNode, ServiceNode], - initial_node: Union[ActiveNode, ServiceNode], - reference_node: Union[ActiveNode, ServiceNode], - config_values: "TrainingConfig", -) -> float: - """ - Calculates score relating to the file system state of a node. - - Args: - final_node: The node after red and blue agents take effect - initial_node: The node before red and blue agents take effect - reference_node: The node if there had been no red or blue effect - """ - score: float = 0.0 - final_node_file_system_state = final_node.file_system_state_actual - reference_node_file_system_state = reference_node.file_system_state_actual - - final_node_scanning_state = final_node.file_system_scanning - reference_node_scanning_state = reference_node.file_system_scanning - - # File System State - if final_node_file_system_state == reference_node_file_system_state: - # All is well - we're no different from the reference situation - score += config_values.all_ok - else: - # We're different from the reference situation - # Need to compare reference and final state of node (i.e. at every step) - if reference_node_file_system_state == FileSystemState.GOOD: - if final_node_file_system_state == FileSystemState.REPAIRING: - score += config_values.repairing_should_be_good - elif final_node_file_system_state == FileSystemState.RESTORING: - score += config_values.restoring_should_be_good - elif final_node_file_system_state == FileSystemState.CORRUPT: - score += config_values.corrupt_should_be_good - elif final_node_file_system_state == FileSystemState.DESTROYED: - score += config_values.destroyed_should_be_good - else: - pass - elif reference_node_file_system_state == FileSystemState.REPAIRING: - if final_node_file_system_state == FileSystemState.GOOD: - score += config_values.good_should_be_repairing - elif final_node_file_system_state == FileSystemState.RESTORING: - score += config_values.restoring_should_be_repairing - elif final_node_file_system_state == FileSystemState.CORRUPT: - score += config_values.corrupt_should_be_repairing - elif final_node_file_system_state == FileSystemState.DESTROYED: - score += config_values.destroyed_should_be_repairing - elif final_node_file_system_state == FileSystemState.REPAIRING: - score += config_values.repairing - else: - pass - elif reference_node_file_system_state == FileSystemState.RESTORING: - if final_node_file_system_state == FileSystemState.GOOD: - score += config_values.good_should_be_restoring - elif final_node_file_system_state == FileSystemState.REPAIRING: - score += config_values.repairing_should_be_restoring - elif final_node_file_system_state == FileSystemState.CORRUPT: - score += config_values.corrupt_should_be_restoring - elif final_node_file_system_state == FileSystemState.DESTROYED: - score += config_values.destroyed_should_be_restoring - elif final_node_file_system_state == FileSystemState.RESTORING: - score += config_values.restoring - else: - pass - elif reference_node_file_system_state == FileSystemState.CORRUPT: - if final_node_file_system_state == FileSystemState.GOOD: - score += config_values.good_should_be_corrupt - elif final_node_file_system_state == FileSystemState.REPAIRING: - score += config_values.repairing_should_be_corrupt - elif final_node_file_system_state == FileSystemState.RESTORING: - score += config_values.restoring_should_be_corrupt - elif final_node_file_system_state == FileSystemState.DESTROYED: - score += config_values.destroyed_should_be_corrupt - elif final_node_file_system_state == FileSystemState.CORRUPT: - score += config_values.corrupt - else: - pass - elif reference_node_file_system_state == FileSystemState.DESTROYED: - if final_node_file_system_state == FileSystemState.GOOD: - score += config_values.good_should_be_destroyed - elif final_node_file_system_state == FileSystemState.REPAIRING: - score += config_values.repairing_should_be_destroyed - elif final_node_file_system_state == FileSystemState.RESTORING: - score += config_values.restoring_should_be_destroyed - elif final_node_file_system_state == FileSystemState.CORRUPT: - score += config_values.corrupt_should_be_destroyed - elif final_node_file_system_state == FileSystemState.DESTROYED: - score += config_values.destroyed - else: - pass - else: - pass - - # Scanning State - if final_node_scanning_state == reference_node_scanning_state: - # All is well - we're no different from the reference situation - score += config_values.all_ok - else: - # We're different from the reference situation - # We're scanning the file system which incurs a penalty (as it slows down systems) - score += config_values.scanning - - return score diff --git a/src/primaite/links/__init__.py b/src/primaite/links/__init__.py deleted file mode 100644 index c91b6951..00000000 --- a/src/primaite/links/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Network connections between nodes in the simulation.""" diff --git a/src/primaite/links/link.py b/src/primaite/links/link.py deleted file mode 100644 index 3830a15b..00000000 --- a/src/primaite/links/link.py +++ /dev/null @@ -1,114 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The link class.""" -from typing import List - -from primaite.common.protocol import Protocol - - -class Link(object): - """Link class.""" - - def __init__(self, _id: str, _bandwidth: int, _source_node_name: str, _dest_node_name: str, _services: str) -> None: - """ - Initialise a Link within the simulated network. - - :param _id: The IER id - :param _bandwidth: The bandwidth of the link (bps) - :param _source_node_name: The name of the source node - :param _dest_node_name: The name of the destination node - :param _protocols: The protocols to add to the link - """ - self.id: str = _id - self.bandwidth: int = _bandwidth - self.source_node_name: str = _source_node_name - self.dest_node_name: str = _dest_node_name - self.protocol_list: List[Protocol] = [] - - # Add the default protocols - for protocol_name in _services: - self.add_protocol(protocol_name) - - def add_protocol(self, _protocol: str) -> None: - """ - Adds a new protocol to the list of protocols on this link. - - Args: - _protocol: The protocol to be added (enum) - """ - self.protocol_list.append(Protocol(_protocol)) - - def get_id(self) -> str: - """ - Gets link ID. - - Returns: - Link ID - """ - return self.id - - def get_source_node_name(self) -> str: - """ - Gets source node name. - - Returns: - Source node name - """ - return self.source_node_name - - def get_dest_node_name(self) -> str: - """ - Gets destination node name. - - Returns: - Destination node name - """ - return self.dest_node_name - - def get_bandwidth(self) -> int: - """ - Gets bandwidth of link. - - Returns: - Link bandwidth (bps) - """ - return self.bandwidth - - def get_protocol_list(self) -> List[Protocol]: - """ - Gets list of protocols on this link. - - Returns: - List of protocols on this link - """ - return self.protocol_list - - def get_current_load(self) -> int: - """ - Gets current total load on this link. - - Returns: - Total load on this link (bps) - """ - total_load = 0 - for protocol in self.protocol_list: - total_load += protocol.get_load() - return total_load - - def add_protocol_load(self, _protocol: str, _load: int) -> None: - """ - Adds a loading to a protocol on this link. - - Args: - _protocol: The protocol to load - _load: The amount to load (bps) - """ - for protocol in self.protocol_list: - if protocol.get_name() == _protocol: - protocol.add_load(_load) - else: - pass - - def clear_traffic(self) -> None: - """Clears all traffic on this link.""" - for protocol in self.protocol_list: - protocol.clear_load() diff --git a/src/primaite/nodes/__init__.py b/src/primaite/nodes/__init__.py deleted file mode 100644 index 231b8d92..00000000 --- a/src/primaite/nodes/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Nodes represent network hosts in the simulation.""" diff --git a/src/primaite/nodes/active_node.py b/src/primaite/nodes/active_node.py deleted file mode 100644 index 8f472e86..00000000 --- a/src/primaite/nodes/active_node.py +++ /dev/null @@ -1,208 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""An Active Node (i.e. not an actuator).""" -import logging -from typing import Final - -from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState -from primaite.config.training_config import TrainingConfig -from primaite.nodes.node import Node - -_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) - - -class ActiveNode(Node): - """Active Node class.""" - - def __init__( - self, - node_id: str, - name: str, - node_type: NodeType, - priority: Priority, - hardware_state: HardwareState, - ip_address: str, - software_state: SoftwareState, - file_system_state: FileSystemState, - config_values: TrainingConfig, - ) -> None: - """ - Initialise an active node. - - :param node_id: The node ID - :param name: The node name - :param node_type: The node type (enum) - :param priority: The node priority (enum) - :param hardware_state: The node Hardware State - :param ip_address: The node IP address - :param software_state: The node Software State - :param file_system_state: The node file system state - :param config_values: The config values - """ - super().__init__(node_id, name, node_type, priority, hardware_state, config_values) - self.ip_address: str = ip_address - # Related to Software - self._software_state: SoftwareState = software_state - self.patching_count: int = 0 - # Related to File System - self.file_system_state_actual: FileSystemState = file_system_state - self.file_system_state_observed: FileSystemState = file_system_state - self.file_system_scanning: bool = False - self.file_system_scanning_count: int = 0 - self.file_system_action_count: int = 0 - - @property - def software_state(self) -> SoftwareState: - """ - Get the software_state. - - :return: The software_state. - """ - return self._software_state - - @software_state.setter - def software_state(self, software_state: SoftwareState) -> None: - """ - Get the software_state. - - :param software_state: Software State. - """ - if self.hardware_state != HardwareState.OFF: - self._software_state = software_state - if software_state == SoftwareState.PATCHING: - self.patching_count = self.config_values.os_patching_duration - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so OS State cannot be " - f"changed. " - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.software_state:{self._software_state}" - ) - - def set_software_state_if_not_compromised(self, software_state: SoftwareState) -> None: - """ - Sets Software State if the node is not compromised. - - Args: - software_state: Software State - """ - if self.hardware_state != HardwareState.OFF: - if self._software_state != SoftwareState.COMPROMISED: - self._software_state = software_state - if software_state == SoftwareState.PATCHING: - self.patching_count = self.config_values.os_patching_duration - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so OS State cannot be changed." - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.software_state:{self._software_state}" - ) - - def update_os_patching_status(self) -> None: - """Updates operating system status based on patching cycle.""" - self.patching_count -= 1 - if self.patching_count <= 0: - self.patching_count = 0 - self._software_state = SoftwareState.GOOD - - def set_file_system_state(self, file_system_state: FileSystemState) -> None: - """ - Sets the file system state (actual and observed). - - Args: - file_system_state: File system state - """ - if self.hardware_state != HardwareState.OFF: - self.file_system_state_actual = file_system_state - - if file_system_state == FileSystemState.REPAIRING: - self.file_system_action_count = self.config_values.file_system_repairing_limit - self.file_system_state_observed = FileSystemState.REPAIRING - elif file_system_state == FileSystemState.RESTORING: - self.file_system_action_count = self.config_values.file_system_restoring_limit - self.file_system_state_observed = FileSystemState.RESTORING - elif file_system_state == FileSystemState.GOOD: - self.file_system_state_observed = FileSystemState.GOOD - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so File System State " - f"cannot be changed. " - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.file_system_state.actual:{self.file_system_state_actual}" - ) - - def set_file_system_state_if_not_compromised(self, file_system_state: FileSystemState) -> None: - """ - Sets the file system state (actual and observed) if not in a compromised state. - - Use for green PoL to prevent it overturning a compromised state - - Args: - file_system_state: File system state - """ - if self.hardware_state != HardwareState.OFF: - if ( - self.file_system_state_actual != FileSystemState.CORRUPT - and self.file_system_state_actual != FileSystemState.DESTROYED - ): - self.file_system_state_actual = file_system_state - - if file_system_state == FileSystemState.REPAIRING: - self.file_system_action_count = self.config_values.file_system_repairing_limit - self.file_system_state_observed = FileSystemState.REPAIRING - elif file_system_state == FileSystemState.RESTORING: - self.file_system_action_count = self.config_values.file_system_restoring_limit - self.file_system_state_observed = FileSystemState.RESTORING - elif file_system_state == FileSystemState.GOOD: - self.file_system_state_observed = FileSystemState.GOOD - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so File System State (if not " - f"compromised) cannot be changed. " - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.file_system_state.actual:{self.file_system_state_actual}" - ) - - def start_file_system_scan(self) -> None: - """Starts a file system scan.""" - self.file_system_scanning = True - self.file_system_scanning_count = self.config_values.file_system_scanning_limit - - def update_file_system_state(self) -> None: - """Updates file system status based on scanning/restore/repair cycle.""" - # Deprecate both the action count (for restoring or reparing) and the scanning count - self.file_system_action_count -= 1 - self.file_system_scanning_count -= 1 - - # Reparing / Restoring updates - if self.file_system_action_count <= 0: - self.file_system_action_count = 0 - if ( - self.file_system_state_actual == FileSystemState.REPAIRING - or self.file_system_state_actual == FileSystemState.RESTORING - ): - self.file_system_state_actual = FileSystemState.GOOD - self.file_system_state_observed = FileSystemState.GOOD - - # Scanning updates - if self.file_system_scanning == True and self.file_system_scanning_count < 0: - self.file_system_state_observed = self.file_system_state_actual - self.file_system_scanning = False - self.file_system_scanning_count = 0 - - def update_resetting_status(self) -> None: - """Updates the reset count & makes software and file state to GOOD.""" - super().update_resetting_status() - if self.resetting_count <= 0: - self.file_system_state_actual = FileSystemState.GOOD - self.software_state = SoftwareState.GOOD - - def update_booting_status(self) -> None: - """Updates the booting software and file state to GOOD.""" - super().update_booting_status() - if self.booting_count <= 0: - self.file_system_state_actual = FileSystemState.GOOD - self.software_state = SoftwareState.GOOD diff --git a/src/primaite/nodes/node.py b/src/primaite/nodes/node.py deleted file mode 100644 index fc4d41d3..00000000 --- a/src/primaite/nodes/node.py +++ /dev/null @@ -1,79 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The base Node class.""" -from typing import Final - -from primaite.common.enums import HardwareState, NodeType, Priority -from primaite.config.training_config import TrainingConfig - - -class Node: - """Node class.""" - - def __init__( - self, - node_id: str, - name: str, - node_type: NodeType, - priority: Priority, - hardware_state: HardwareState, - config_values: TrainingConfig, - ) -> None: - """ - Initialise a node. - - :param node_id: The node id. - :param name: The name of the node. - :param node_type: The type of the node. - :param priority: The priority of the node. - :param hardware_state: The state of the node. - :param config_values: Config values. - """ - self.node_id: Final[str] = node_id - self.name: Final[str] = name - self.node_type: Final[NodeType] = node_type - self.priority = priority - self.hardware_state: HardwareState = hardware_state - self.resetting_count: int = 0 - self.config_values: TrainingConfig = config_values - self.booting_count: int = 0 - self.shutting_down_count: int = 0 - - def __repr__(self) -> str: - """Returns the name of the node.""" - return self.name - - def turn_on(self) -> None: - """Sets the node state to ON.""" - self.hardware_state = HardwareState.BOOTING - self.booting_count = self.config_values.node_booting_duration - - def turn_off(self) -> None: - """Sets the node state to OFF.""" - self.hardware_state = HardwareState.OFF - self.shutting_down_count = self.config_values.node_shutdown_duration - - def reset(self) -> None: - """Sets the node state to Resetting and starts the reset count.""" - self.hardware_state = HardwareState.RESETTING - self.resetting_count = self.config_values.node_reset_duration - - def update_resetting_status(self) -> None: - """Updates the resetting count.""" - self.resetting_count -= 1 - if self.resetting_count <= 0: - self.resetting_count = 0 - self.hardware_state = HardwareState.ON - - def update_booting_status(self) -> None: - """Updates the booting count.""" - self.booting_count -= 1 - if self.booting_count <= 0: - self.booting_count = 0 - self.hardware_state = HardwareState.ON - - def update_shutdown_status(self) -> None: - """Updates the shutdown count.""" - self.shutting_down_count -= 1 - if self.shutting_down_count <= 0: - self.shutting_down_count = 0 - self.hardware_state = HardwareState.OFF diff --git a/src/primaite/nodes/node_state_instruction_green.py b/src/primaite/nodes/node_state_instruction_green.py deleted file mode 100644 index 6e35d0ec..00000000 --- a/src/primaite/nodes/node_state_instruction_green.py +++ /dev/null @@ -1,94 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Defines node behaviour for Green PoL.""" -from typing import TYPE_CHECKING, Union - -if TYPE_CHECKING: - from primaite.common.enums import FileSystemState, HardwareState, NodePOLType, SoftwareState - - -class NodeStateInstructionGreen(object): - """The Node State Instruction class.""" - - def __init__( - self, - _id: str, - _start_step: int, - _end_step: int, - _node_id: str, - _node_pol_type: "NodePOLType", - _service_name: str, - _state: Union["HardwareState", "SoftwareState", "FileSystemState"], - ) -> None: - """ - Initialise the Node State Instruction. - - :param _id: The node state instruction id - :param _start_step: The start step of the instruction - :param _end_step: The end step of the instruction - :param _node_id: The id of the associated node - :param _node_pol_type: The pattern of life type - :param _service_name: The service name - :param _state: The state (node or service) - """ - self.id = _id - self.start_step = _start_step - self.end_step = _end_step - self.node_id = _node_id - self.node_pol_type: "NodePOLType" = _node_pol_type - self.service_name: str = _service_name # Not used when not a service instruction - # TODO: confirm type of state - self.state: Union["HardwareState", "SoftwareState", "FileSystemState"] = _state - - def get_start_step(self) -> int: - """ - Gets the start step. - - Returns: - The start step - """ - return self.start_step - - def get_end_step(self) -> int: - """ - Gets the end step. - - Returns: - The end step - """ - return self.end_step - - def get_node_id(self) -> str: - """ - Gets the node ID. - - Returns: - The node ID - """ - return self.node_id - - def get_node_pol_type(self) -> "NodePOLType": - """ - Gets the node pattern of life type (enum). - - Returns: - The node pattern of life type (enum) - """ - return self.node_pol_type - - def get_service_name(self) -> str: - """ - Gets the service name. - - Returns: - The service name - """ - return self.service_name - - def get_state(self) -> Union["HardwareState", "SoftwareState", "FileSystemState"]: - """ - Gets the state (node or service). - - Returns: - The state (node or service) - """ - return self.state diff --git a/src/primaite/nodes/node_state_instruction_red.py b/src/primaite/nodes/node_state_instruction_red.py deleted file mode 100644 index eb87924b..00000000 --- a/src/primaite/nodes/node_state_instruction_red.py +++ /dev/null @@ -1,143 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Defines node behaviour for Green PoL.""" -from typing import TYPE_CHECKING, Union - -from primaite.common.enums import NodePOLType - -if TYPE_CHECKING: - from primaite.common.enums import FileSystemState, HardwareState, NodePOLInitiator, SoftwareState - - -class NodeStateInstructionRed: - """The Node State Instruction class.""" - - def __init__( - self, - _id: str, - _start_step: int, - _end_step: int, - _target_node_id: str, - _pol_initiator: "NodePOLInitiator", - _pol_type: NodePOLType, - pol_protocol: str, - _pol_state: Union["HardwareState", "SoftwareState", "FileSystemState"], - _pol_source_node_id: str, - _pol_source_node_service: str, - _pol_source_node_service_state: str, - ) -> None: - """ - Initialise the Node State Instruction for the red agent. - - :param _id: The node state instruction id - :param _start_step: The start step of the instruction - :param _end_step: The end step of the instruction - :param _target_node_id: The id of the associated node - :param -pol_initiator: The way the PoL is applied (DIRECT, IER or SERVICE) - :param _pol_type: The pattern of life type - :param pol_protocol: The pattern of life protocol/service affected - :param _pol_state: The state (node or service) - :param _pol_source_node_id: The source node Id (used for initiator type SERVICE) - :param _pol_source_node_service: The source node service (used for initiator type SERVICE) - :param _pol_source_node_service_state: The source node service state (used for initiator type SERVICE) - """ - self.id: str = _id - self.start_step: int = _start_step - self.end_step: int = _end_step - self.target_node_id: str = _target_node_id - self.initiator: "NodePOLInitiator" = _pol_initiator - self.pol_type: NodePOLType = _pol_type - self.service_name: str = pol_protocol # Not used when not a service instruction - self.state: Union["HardwareState", "SoftwareState", "FileSystemState"] = _pol_state - self.source_node_id: str = _pol_source_node_id - self.source_node_service: str = _pol_source_node_service - self.source_node_service_state = _pol_source_node_service_state - - def get_start_step(self) -> int: - """ - Gets the start step. - - Returns: - The start step - """ - return self.start_step - - def get_end_step(self) -> int: - """ - Gets the end step. - - Returns: - The end step - """ - return self.end_step - - def get_target_node_id(self) -> str: - """ - Gets the node ID. - - Returns: - The node ID - """ - return self.target_node_id - - def get_initiator(self) -> "NodePOLInitiator": - """ - Gets the initiator. - - Returns: - The initiator - """ - return self.initiator - - def get_pol_type(self) -> NodePOLType: - """ - Gets the node pattern of life type (enum). - - Returns: - The node pattern of life type (enum) - """ - return self.pol_type - - def get_service_name(self) -> str: - """ - Gets the service name. - - Returns: - The service name - """ - return self.service_name - - def get_state(self) -> Union["HardwareState", "SoftwareState", "FileSystemState"]: - """ - Gets the state (node or service). - - Returns: - The state (node or service) - """ - return self.state - - def get_source_node_id(self) -> str: - """ - Gets the source node id (used for initiator type SERVICE). - - Returns: - The source node id - """ - return self.source_node_id - - def get_source_node_service(self) -> str: - """ - Gets the source node service (used for initiator type SERVICE). - - Returns: - The source node service - """ - return self.source_node_service - - def get_source_node_service_state(self) -> str: - """ - Gets the source node service state (used for initiator type SERVICE). - - Returns: - The source node service state - """ - return self.source_node_service_state diff --git a/src/primaite/nodes/passive_node.py b/src/primaite/nodes/passive_node.py deleted file mode 100644 index 08dcbfa2..00000000 --- a/src/primaite/nodes/passive_node.py +++ /dev/null @@ -1,42 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The Passive Node class (i.e. an actuator).""" -from primaite.common.enums import HardwareState, NodeType, Priority -from primaite.config.training_config import TrainingConfig -from primaite.nodes.node import Node - - -class PassiveNode(Node): - """The Passive Node class.""" - - def __init__( - self, - node_id: str, - name: str, - node_type: NodeType, - priority: Priority, - hardware_state: HardwareState, - config_values: TrainingConfig, - ) -> None: - """ - Initialise a passive node. - - :param node_id: The node id. - :param name: The name of the node. - :param node_type: The type of the node. - :param priority: The priority of the node. - :param hardware_state: The state of the node. - :param config_values: Config values. - """ - # Pass through to Super for now - super().__init__(node_id, name, node_type, priority, hardware_state, config_values) - - @property - def ip_address(self) -> str: - """ - Gets the node IP address as an empty string. - - No concept of IP address for passive nodes for now. - - :return: The node IP address. - """ - return "" diff --git a/src/primaite/nodes/service_node.py b/src/primaite/nodes/service_node.py deleted file mode 100644 index b0d42785..00000000 --- a/src/primaite/nodes/service_node.py +++ /dev/null @@ -1,190 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""A Service Node (i.e. not an actuator).""" -import logging -from typing import Dict, Final - -from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState -from primaite.common.service import Service -from primaite.config.training_config import TrainingConfig -from primaite.nodes.active_node import ActiveNode - -_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) - - -class ServiceNode(ActiveNode): - """ServiceNode class.""" - - def __init__( - self, - node_id: str, - name: str, - node_type: NodeType, - priority: Priority, - hardware_state: HardwareState, - ip_address: str, - software_state: SoftwareState, - file_system_state: FileSystemState, - config_values: TrainingConfig, - ) -> None: - """ - Initialise a Service Node. - - :param node_id: The node ID - :param name: The node name - :param node_type: The node type (enum) - :param priority: The node priority (enum) - :param hardware_state: The node Hardware State - :param ip_address: The node IP address - :param software_state: The node Software State - :param file_system_state: The node file system state - :param config_values: The config values - """ - super().__init__( - node_id, - name, - node_type, - priority, - hardware_state, - ip_address, - software_state, - file_system_state, - config_values, - ) - self.services: Dict[str, Service] = {} - - def add_service(self, service: Service) -> None: - """ - Adds a service to the node. - - :param service: The service to add - """ - self.services[service.name] = service - - def has_service(self, protocol_name: str) -> bool: - """ - Indicates whether a service is on a node. - - :param protocol_name: The service (protocol)e. - :return: True if service (protocol) is on the node, otherwise False. - """ - for service_key, service_value in self.services.items(): - if service_key == protocol_name: - return True - return False - - def service_running(self, protocol_name: str) -> bool: - """ - Indicates whether a service is in a running state on the node. - - :param protocol_name: The service (protocol) - :return: True if service (protocol) is in a running state on the node, otherwise False. - """ - for service_key, service_value in self.services.items(): - if service_key == protocol_name: - if service_value.software_state != SoftwareState.PATCHING: - return True - else: - return False - return False - - def service_is_overwhelmed(self, protocol_name: str) -> bool: - """ - Indicates whether a service is in an overwhelmed state on the node. - - :param protocol_name: The service (protocol) - :return: True if service (protocol) is in an overwhelmed state on the node, otherwise False. - """ - for service_key, service_value in self.services.items(): - if service_key == protocol_name: - if service_value.software_state == SoftwareState.OVERWHELMED: - return True - else: - return False - return False - - def set_service_state(self, protocol_name: str, software_state: SoftwareState) -> None: - """ - Sets the software_state of a service (protocol) on the node. - - :param protocol_name: The service (protocol). - :param software_state: The software_state. - """ - if self.hardware_state != HardwareState.OFF: - service_key = protocol_name - service_value = self.services.get(service_key) - if service_value: - # Can't set to compromised if you're in a patching state - if ( - software_state == SoftwareState.COMPROMISED - and service_value.software_state != SoftwareState.PATCHING - ) or software_state != SoftwareState.COMPROMISED: - service_value.software_state = software_state - if software_state == SoftwareState.PATCHING: - service_value.patching_count = self.config_values.service_patching_duration - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so the state of a service " - f"cannot be changed. " - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.services[]:{protocol_name}, " - f"Node.services[].software_state:{software_state}" - ) - - def set_service_state_if_not_compromised(self, protocol_name: str, software_state: SoftwareState) -> None: - """ - Sets the software_state of a service (protocol) on the node. - - Done if the software_state is not "compromised". - - :param protocol_name: The service (protocol). - :param software_state: The software_state. - """ - if self.hardware_state != HardwareState.OFF: - service_key = protocol_name - service_value = self.services.get(service_key) - if service_value: - if service_value.software_state != SoftwareState.COMPROMISED: - service_value.software_state = software_state - if software_state == SoftwareState.PATCHING: - service_value.patching_count = self.config_values.service_patching_duration - else: - _LOGGER.info( - f"The Nodes hardware state is OFF so the state of a service " - f"cannot be changed. " - f"Node.node_id:{self.node_id}, " - f"Node.hardware_state:{self.hardware_state}, " - f"Node.services[]:{protocol_name}, " - f"Node.services[].software_state:{software_state}" - ) - - def get_service_state(self, protocol_name: str) -> SoftwareState: - """ - Gets the state of a service. - - :return: The software_state of the service. - """ - service_key = protocol_name - service_value = self.services.get(service_key) - if service_value: - return service_value.software_state - - def update_services_patching_status(self) -> None: - """Updates the patching counter for any service that are patching.""" - for service_key, service_value in self.services.items(): - if service_value.software_state == SoftwareState.PATCHING: - service_value.reduce_patching_count() - - def update_resetting_status(self) -> None: - """Update resetting counter and set software state if it reached 0.""" - super().update_resetting_status() - if self.resetting_count <= 0: - for service in self.services.values(): - service.software_state = SoftwareState.GOOD - - def update_booting_status(self) -> None: - """Update booting counter and set software to good if it reached 0.""" - super().update_booting_status() - if self.booting_count <= 0: - for service in self.services.values(): - service.software_state = SoftwareState.GOOD diff --git a/src/primaite/pol/__init__.py b/src/primaite/pol/__init__.py deleted file mode 100644 index d0d9f616..00000000 --- a/src/primaite/pol/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Pattern of Life- Represents the actions of users on the network.""" diff --git a/src/primaite/pol/green_pol.py b/src/primaite/pol/green_pol.py deleted file mode 100644 index 814aa314..00000000 --- a/src/primaite/pol/green_pol.py +++ /dev/null @@ -1,264 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Implements Pattern of Life on the network (nodes and links).""" -from typing import Dict - -from networkx import MultiGraph, shortest_path - -from primaite.acl.access_control_list import AccessControlList -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import HardwareState, NodePOLType, NodeType, SoftwareState -from primaite.links.link import Link -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen -from primaite.nodes.service_node import ServiceNode -from primaite.pol.ier import IER - -_VERBOSE: bool = False - - -def apply_iers( - network: MultiGraph, - nodes: Dict[str, NodeUnion], - links: Dict[str, Link], - iers: Dict[str, IER], - acl: AccessControlList, - step: int, -) -> None: - """ - Applies IERs to the links (link pattern of life). - - Args: - network: The network modelled in the environment - nodes: The nodes within the environment - links: The links within the environment - iers: The IERs to apply to the links - acl: The Access Control List - step: The step number. - """ - if _VERBOSE: - print("Applying IERs") - - # Go through each IER and check the conditions for it being applied - # If everything is in place, apply the IER protocol load to the relevant links - for ier_key, ier_value in iers.items(): - start_step = ier_value.get_start_step() - stop_step = ier_value.get_end_step() - protocol = ier_value.get_protocol() - port = ier_value.get_port() - load = ier_value.get_load() - source_node_id = ier_value.get_source_node_id() - dest_node_id = ier_value.get_dest_node_id() - - # Need to set the running status to false first for all IERs - ier_value.set_is_running(False) - - source_valid = True - dest_valid = True - acl_block = False - - if step >= start_step and step <= stop_step: - # continue -------------------------- - - # Get the source and destination node for this link - source_node = nodes[source_node_id] - dest_node = nodes[dest_node_id] - - # 1. Check the source node situation - # TODO: should be using isinstance rather than checking node type attribute. IE. just because it's a switch - # doesn't mean it has a software state? It could be a PassiveNode or ActiveNode - if source_node.node_type == NodeType.SWITCH: - # It's a switch - if ( - source_node.hardware_state == HardwareState.ON - and source_node.software_state != SoftwareState.PATCHING - ): - source_valid = True - else: - # IER no longer valid - source_valid = False - elif source_node.node_type == NodeType.ACTUATOR: - # It's an actuator - # TO DO - pass - else: - # It's not a switch or an actuator (so active node) - if ( - source_node.hardware_state == HardwareState.ON - and source_node.software_state != SoftwareState.PATCHING - ): - if source_node.has_service(protocol): - if source_node.service_running(protocol) and not source_node.service_is_overwhelmed(protocol): - source_valid = True - else: - source_valid = False - else: - # Do nothing - IER is not valid on this node - # (This shouldn't happen if the IER has been written correctly) - source_valid = False - else: - # Do nothing - IER no longer valid - source_valid = False - - # 2. Check the dest node situation - if dest_node.node_type == NodeType.SWITCH: - # It's a switch - if dest_node.hardware_state == HardwareState.ON and dest_node.software_state != SoftwareState.PATCHING: - dest_valid = True - else: - # IER no longer valid - dest_valid = False - elif dest_node.node_type == NodeType.ACTUATOR: - # It's an actuator - pass - else: - # It's not a switch or an actuator (so active node) - if dest_node.hardware_state == HardwareState.ON and dest_node.software_state != SoftwareState.PATCHING: - if dest_node.has_service(protocol): - if dest_node.service_running(protocol) and not dest_node.service_is_overwhelmed(protocol): - dest_valid = True - else: - dest_valid = False - else: - # Do nothing - IER is not valid on this node - # (This shouldn't happen if the IER has been written correctly) - dest_valid = False - else: - # Do nothing - IER no longer valid - dest_valid = False - - # 3. Check that the ACL doesn't block it - acl_block = acl.is_blocked(source_node.ip_address, dest_node.ip_address, protocol, port) - if acl_block: - if _VERBOSE: - print( - "ACL block on source: " - + source_node.ip_address - + ", dest: " - + dest_node.ip_address - + ", protocol: " - + protocol - + ", port: " - + port - ) - else: - if _VERBOSE: - print("No ACL block") - - # Check whether both the source and destination are valid, and there's no ACL block - if source_valid and dest_valid and not acl_block: - # Load up the link(s) with the traffic - - if _VERBOSE: - print("Source, Dest and ACL valid") - - # Get the shortest path (i.e. nodes) between source and destination - path_node_list = shortest_path(network, source_node, dest_node) - path_node_list_length = len(path_node_list) - path_valid = True - - # We might have a switch in the path, so check all nodes are operational - for node in path_node_list: - if node.hardware_state != HardwareState.ON or node.software_state == SoftwareState.PATCHING: - path_valid = False - - if path_valid: - if _VERBOSE: - print("Applying IER to link(s)") - count = 0 - link_capacity_exceeded = False - - # Check that the link capacity is not exceeded by the new load - while count < path_node_list_length - 1: - # Get the link between the next two nodes - edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count + 1]) - link_id = edge_dict[0].get("id") - link = links[link_id] - # Check whether the new load exceeds the bandwidth - if (link.get_current_load() + load) > link.get_bandwidth(): - link_capacity_exceeded = True - if _VERBOSE: - print("Link capacity exceeded") - pass - count += 1 - - # Check whether the link capacity for any links on this path have been exceeded - if link_capacity_exceeded == False: - # Now apply the new loads to the links - count = 0 - while count < path_node_list_length - 1: - # Get the link between the next two nodes - edge_dict = network.get_edge_data( - path_node_list[count], - path_node_list[count + 1], - ) - link_id = edge_dict[0].get("id") - link = links[link_id] - # Add the load from this IER - link.add_protocol_load(protocol, load) - count += 1 - # This IER is now valid, so set it to running - ier_value.set_is_running(True) - else: - # One of the nodes is not operational - if _VERBOSE: - print("Path not valid - one or more nodes not operational") - pass - - else: - if _VERBOSE: - print("Source, Dest or ACL were not valid") - pass - # ------------------------------------ - else: - # Do nothing - IER no longer valid - pass - - -def apply_node_pol( - nodes: Dict[str, NodeUnion], - node_pol: Dict[str, NodeStateInstructionGreen], - step: int, -) -> None: - """ - Applies node pattern of life. - - Args: - nodes: The nodes within the environment - node_pol: The node pattern of life to apply - step: The step number. - """ - if _VERBOSE: - print("Applying Node PoL") - - for key, node_instruction in node_pol.items(): - start_step = node_instruction.get_start_step() - stop_step = node_instruction.get_end_step() - node_id = node_instruction.get_node_id() - node_pol_type = node_instruction.get_node_pol_type() - service_name = node_instruction.get_service_name() - state = node_instruction.get_state() - - if step >= start_step and step <= stop_step: - # continue -------------------------- - node = nodes[node_id] - - if node_pol_type == NodePOLType.OPERATING: - # Change hardware state - node.hardware_state = state - elif node_pol_type == NodePOLType.OS: - # Change OS state - # Don't allow PoL to fix something that is compromised. Only the Blue agent can do this - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - node.set_software_state_if_not_compromised(state) - elif node_pol_type == NodePOLType.SERVICE: - # Change a service state - # Don't allow PoL to fix something that is compromised. Only the Blue agent can do this - if isinstance(node, ServiceNode): - node.set_service_state_if_not_compromised(service_name, state) - else: - # Change the file system status - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - node.set_file_system_state_if_not_compromised(state) - else: - # PoL is not valid in this time step - pass diff --git a/src/primaite/pol/ier.py b/src/primaite/pol/ier.py deleted file mode 100644 index b8da28c0..00000000 --- a/src/primaite/pol/ier.py +++ /dev/null @@ -1,147 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -""" -Information Exchange Requirements for APE. - -Used to represent an information flow from source to destination. -""" - - -class IER(object): - """Information Exchange Requirement class.""" - - def __init__( - self, - _id: str, - _start_step: int, - _end_step: int, - _load: int, - _protocol: str, - _port: str, - _source_node_id: str, - _dest_node_id: str, - _mission_criticality: int, - _running: bool = False, - ) -> None: - """ - Initialise an Information Exchange Request. - - :param _id: The IER id - :param _start_step: The step when this IER should start - :param _end_step: The step when this IER should end - :param _load: The load this IER should put on a link (bps) - :param _protocol: The protocol of this IER - :param _port: The port this IER runs on - :param _source_node_id: The source node ID - :param _dest_node_id: The destination node ID - :param _mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical) - :param _running: Indicates whether the IER is currently running - """ - self.id: str = _id - self.start_step: int = _start_step - self.end_step: int = _end_step - self.source_node_id: str = _source_node_id - self.dest_node_id: str = _dest_node_id - self.load: int = _load - self.protocol: str = _protocol - self.port: str = _port - self.mission_criticality: int = _mission_criticality - self.running: bool = _running - - def get_id(self) -> str: - """ - Gets IER ID. - - Returns: - IER ID - """ - return self.id - - def get_start_step(self) -> int: - """ - Gets IER start step. - - Returns: - IER start step - """ - return self.start_step - - def get_end_step(self) -> int: - """ - Gets IER end step. - - Returns: - IER end step - """ - return self.end_step - - def get_load(self) -> int: - """ - Gets IER load. - - Returns: - IER load - """ - return self.load - - def get_protocol(self) -> str: - """ - Gets IER protocol. - - Returns: - IER protocol - """ - return self.protocol - - def get_port(self) -> str: - """ - Gets IER port. - - Returns: - IER port - """ - return self.port - - def get_source_node_id(self) -> str: - """ - Gets IER source node ID. - - Returns: - IER source node ID - """ - return self.source_node_id - - def get_dest_node_id(self) -> str: - """ - Gets IER destination node ID. - - Returns: - IER destination node ID - """ - return self.dest_node_id - - def get_is_running(self) -> bool: - """ - Informs whether the IER is currently running. - - Returns: - True if running - """ - return self.running - - def set_is_running(self, _value: bool) -> None: - """ - Sets the running state of the IER. - - Args: - _value: running status - """ - self.running = _value - - def get_mission_criticality(self) -> int: - """ - Gets the IER mission criticality (used in the reward function). - - Returns: - Mission criticality value (0 lowest to 5 highest) - """ - return self.mission_criticality diff --git a/src/primaite/pol/red_agent_pol.py b/src/primaite/pol/red_agent_pol.py deleted file mode 100644 index ca1a58da..00000000 --- a/src/primaite/pol/red_agent_pol.py +++ /dev/null @@ -1,353 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Implements POL on the network (nodes and links) resulting from the red agent attack.""" -from typing import Dict - -from networkx import MultiGraph, shortest_path - -from primaite import getLogger -from primaite.acl.access_control_list import AccessControlList -from primaite.common.custom_typing import NodeUnion -from primaite.common.enums import HardwareState, NodePOLInitiator, NodePOLType, NodeType, SoftwareState -from primaite.links.link import Link -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed -from primaite.nodes.service_node import ServiceNode -from primaite.pol.ier import IER - -_LOGGER = getLogger(__name__) - -_VERBOSE: bool = False - - -def apply_red_agent_iers( - network: MultiGraph, - nodes: Dict[str, NodeUnion], - links: Dict[str, Link], - iers: Dict[str, IER], - acl: AccessControlList, - step: int, -) -> None: - """ - Applies IERs to the links (link POL) resulting from red agent attack. - - Args: - network: The network modelled in the environment - nodes: The nodes within the environment - links: The links within the environment - iers: The red agent IERs to apply to the links - acl: The Access Control List - step: The step number. - """ - # Go through each IER and check the conditions for it being applied - # If everything is in place, apply the IER protocol load to the relevant links - for ier_key, ier_value in iers.items(): - start_step = ier_value.get_start_step() - stop_step = ier_value.get_end_step() - protocol = ier_value.get_protocol() - port = ier_value.get_port() - load = ier_value.get_load() - source_node_id = ier_value.get_source_node_id() - dest_node_id = ier_value.get_dest_node_id() - - # Need to set the running status to false first for all IERs - ier_value.set_is_running(False) - - source_valid = True - dest_valid = True - acl_block = False - - if step >= start_step and step <= stop_step: - # continue -------------------------- - - # Get the source and destination node for this link - source_node = nodes[source_node_id] - dest_node = nodes[dest_node_id] - - # 1. Check the source node situation - if source_node.node_type == NodeType.SWITCH: - # It's a switch - if source_node.hardware_state == HardwareState.ON: - source_valid = True - else: - # IER no longer valid - source_valid = False - elif source_node.node_type == NodeType.ACTUATOR: - # It's an actuator - # TO DO - pass - else: - # It's not a switch or an actuator (so active node) - # TODO: this occurs after ruling out the possibility that the node is a switch or an actuator, but it - # could still be a passive/active node, therefore it won't have a hardware_state. The logic here needs - # to change according to duck typing. - if source_node.hardware_state == HardwareState.ON: - if source_node.has_service(protocol): - # Red agents IERs can only be valid if the source service is in a compromised state - if source_node.get_service_state(protocol) == SoftwareState.COMPROMISED: - source_valid = True - else: - source_valid = False - else: - # Do nothing - IER is not valid on this node - # (This shouldn't happen if the IER has been written correctly) - source_valid = False - else: - # Do nothing - IER no longer valid - source_valid = False - - # 2. Check the dest node situation - if dest_node.node_type == NodeType.SWITCH: - # It's a switch - if dest_node.hardware_state == HardwareState.ON: - dest_valid = True - else: - # IER no longer valid - dest_valid = False - elif dest_node.node_type == NodeType.ACTUATOR: - # It's an actuator - pass - else: - # It's not a switch or an actuator (so active node) - if dest_node.hardware_state == HardwareState.ON: - if dest_node.has_service(protocol): - # We don't care what state the destination service is in for an IER - dest_valid = True - else: - # Do nothing - IER is not valid on this node - # (This shouldn't happen if the IER has been written correctly) - dest_valid = False - else: - # Do nothing - IER no longer valid - dest_valid = False - - # 3. Check that the ACL doesn't block it - acl_block = acl.is_blocked(source_node.ip_address, dest_node.ip_address, protocol, port) - if acl_block: - if _VERBOSE: - print( - "ACL block on source: " - + source_node.ip_address - + ", dest: " - + dest_node.ip_address - + ", protocol: " - + protocol - + ", port: " - + port - ) - else: - if _VERBOSE: - print("No ACL block") - - # Check whether both the source and destination are valid, and there's no ACL block - if source_valid and dest_valid and not acl_block: - # Load up the link(s) with the traffic - - if _VERBOSE: - print("Source, Dest and ACL valid") - - # Get the shortest path (i.e. nodes) between source and destination - path_node_list = shortest_path(network, source_node, dest_node) - path_node_list_length = len(path_node_list) - path_valid = True - - # We might have a switch in the path, so check all nodes are operational - # We're assuming here that red agents can get past switches that are patching - for node in path_node_list: - if node.hardware_state != HardwareState.ON: - path_valid = False - - if path_valid: - if _VERBOSE: - print("Applying IER to link(s)") - count = 0 - link_capacity_exceeded = False - - # Check that the link capacity is not exceeded by the new load - while count < path_node_list_length - 1: - # Get the link between the next two nodes - edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count + 1]) - link_id = edge_dict[0].get("id") - link = links[link_id] - # Check whether the new load exceeds the bandwidth - if (link.get_current_load() + load) > link.get_bandwidth(): - link_capacity_exceeded = True - if _VERBOSE: - print("Link capacity exceeded") - pass - count += 1 - - # Check whether the link capacity for any links on this path have been exceeded - if link_capacity_exceeded == False: - # Now apply the new loads to the links - count = 0 - while count < path_node_list_length - 1: - # Get the link between the next two nodes - edge_dict = network.get_edge_data( - path_node_list[count], - path_node_list[count + 1], - ) - link_id = edge_dict[0].get("id") - link = links[link_id] - # Add the load from this IER - link.add_protocol_load(protocol, load) - count += 1 - # This IER is now valid, so set it to running - ier_value.set_is_running(True) - if _VERBOSE: - print("Red IER was allowed to run in step " + str(step)) - else: - # One of the nodes is not operational - if _VERBOSE: - print("Path not valid - one or more nodes not operational") - pass - - else: - if _VERBOSE: - print("Red IER was NOT allowed to run in step " + str(step)) - print("Source, Dest or ACL were not valid") - pass - # ------------------------------------ - else: - # Do nothing - IER no longer valid - pass - - pass - - -def apply_red_agent_node_pol( - nodes: Dict[str, NodeUnion], - iers: Dict[str, IER], - node_pol: Dict[str, NodeStateInstructionRed], - step: int, -) -> None: - """ - Applies node pattern of life. - - Args: - nodes: The nodes within the environment - iers: The red agent IERs - node_pol: The red agent node pattern of life to apply - step: The step number. - """ - if _VERBOSE: - print("Applying Node Red Agent PoL") - - for key, node_instruction in node_pol.items(): - start_step = node_instruction.get_start_step() - stop_step = node_instruction.get_end_step() - target_node_id = node_instruction.get_target_node_id() - initiator = node_instruction.get_initiator() - pol_type = node_instruction.get_pol_type() - service_name = node_instruction.get_service_name() - state = node_instruction.get_state() - source_node_id = node_instruction.get_source_node_id() - source_node_service_name = node_instruction.get_source_node_service() - source_node_service_state_value = node_instruction.get_source_node_service_state() - - passed_checks = False - - if step >= start_step and step <= stop_step: - # continue -------------------------- - target_node: NodeUnion = nodes[target_node_id] - - # check if the initiator type is a str, and if so, cast it as - # NodePOLInitiator - if isinstance(initiator, str): - initiator = NodePOLInitiator[initiator] - - # Based the action taken on the initiator type - if initiator == NodePOLInitiator.DIRECT: - # No conditions required, just apply the change - passed_checks = True - elif initiator == NodePOLInitiator.IER: - # Need to check there is a red IER incoming - passed_checks = is_red_ier_incoming(target_node, iers, pol_type) - elif initiator == NodePOLInitiator.SERVICE: - # Need to check the condition of a service on another node - source_node = nodes[source_node_id] - if source_node.has_service(source_node_service_name): - if ( - source_node.get_service_state(source_node_service_name) - == SoftwareState[source_node_service_state_value] - ): - passed_checks = True - else: - # Do nothing, no matching state value - pass - else: - # Do nothing, service not on this node - pass - else: - _LOGGER.warning("Node Red Agent PoL not allowed - misconfiguration") - - # Only apply the PoL if the checks have passed (based on the initiator type) - if passed_checks: - # Apply the change - if pol_type == NodePOLType.OPERATING: - # Change hardware state - target_node.hardware_state = state - elif pol_type == NodePOLType.OS: - # Change OS state - if isinstance(target_node, ActiveNode) or isinstance(target_node, ServiceNode): - target_node.software_state = state - elif pol_type == NodePOLType.SERVICE: - # Change a service state - if isinstance(target_node, ServiceNode): - target_node.set_service_state(service_name, state) - else: - # Change the file system status - if isinstance(target_node, ActiveNode) or isinstance(target_node, ServiceNode): - target_node.set_file_system_state(state) - else: - _LOGGER.debug("Node Red Agent PoL not allowed - did not pass checks") - else: - # PoL is not valid in this time step - pass - - -def is_red_ier_incoming(node: NodeUnion, iers: Dict[str, IER], node_pol_type: NodePOLType) -> bool: - """Checks if the RED IER is incoming. - - :param node: Destination node of the IER - :type node: NodeUnion - :param iers: Directory of IERs - :type iers: Dict[str,IER] - :param node_pol_type: Type of Pattern-Of-Life - :type node_pol_type: NodePOLType - :return: Whether the RED IER is incoming. - :rtype: bool - """ - node_id = node.node_id - - for ier_key, ier_value in iers.items(): - if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id: - if ( - node_pol_type == NodePOLType.OPERATING - or node_pol_type == NodePOLType.OS - or node_pol_type == NodePOLType.FILE - ): - # It's looking to change hardware state, file system or SoftwareState, so valid - return True - elif node_pol_type == NodePOLType.SERVICE: - # Check if the service is present on the node and running - ier_protocol = ier_value.get_protocol() - if isinstance(node, ServiceNode): - if node.has_service(ier_protocol): - if node.service_running(ier_protocol): - # Matching service is present and running, so valid - return True - else: - # Service is present, but not running - return False - else: - # Service is not present - return False - else: - # Not a service node - return False - else: - # Shouldn't get here - instruction type is undefined - return False - else: - # The IER destination is not this node, or the IER is not running - return False diff --git a/src/primaite/primaite_session.py b/src/primaite/primaite_session.py deleted file mode 100644 index 7d5b2709..00000000 --- a/src/primaite/primaite_session.py +++ /dev/null @@ -1,228 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Main entry point to PrimAITE. Configure training/evaluation experiments and input/output.""" -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, Final, Optional, Tuple, Union - -from primaite import getLogger -from primaite.agents.agent_abc import AgentSessionABC -from primaite.agents.hardcoded_acl import HardCodedACLAgent -from primaite.agents.hardcoded_node import HardCodedNodeAgent - -# from primaite.agents.rllib import RLlibAgent -from primaite.agents.sb3 import SB3Agent -from primaite.agents.simple import DoNothingACLAgent, DoNothingNodeAgent, DummyAgent, RandomAgent -from primaite.common.enums import ActionType, AgentFramework, AgentIdentifier, SessionType -from primaite.config import lay_down_config, training_config -from primaite.config.training_config import TrainingConfig -from primaite.utils.session_metadata_parser import parse_session_metadata -from primaite.utils.session_output_reader import all_transactions_dict, av_rewards_dict - -_LOGGER = getLogger(__name__) - - -class PrimaiteSession: - """ - The PrimaiteSession class. - - Provides a single learning and evaluation entry point for all training and lay down configurations. - """ - - def __init__( - self, - training_config_path: Optional[Union[str, Path]] = "", - lay_down_config_path: Optional[Union[str, Path]] = "", - session_path: Optional[Union[str, Path]] = None, - legacy_training_config: bool = False, - legacy_lay_down_config: bool = False, - ) -> None: - """ - The PrimaiteSession constructor. - - :param training_config_path: YAML file containing configurable items defined in - `primaite.config.training_config.TrainingConfig` - :type training_config_path: Union[path, str] - :param lay_down_config_path: YAML file containing configurable items for generating network laydown. - :type lay_down_config_path: Union[path, str] - :param session_path: directory path of the session to load - :param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0, - otherwise False. - :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, - otherwise False. - """ - self._agent_session: AgentSessionABC = None # noqa - self.session_path: Path = session_path # noqa - self.timestamp_str: str = None # noqa - self.learning_path: Path = None # noqa - self.evaluation_path: Path = None # noqa - self.legacy_training_config = legacy_training_config - self.legacy_lay_down_config = legacy_lay_down_config - - # check if session path is provided - if session_path is not None: - # set load_session to true - self.is_load_session = True - if not isinstance(session_path, Path): - session_path = Path(session_path) - - # if a session path is provided, load it - if not session_path.exists(): - raise Exception(f"Session could not be loaded. Path does not exist: {session_path}") - - md_dict, training_config_path, lay_down_config_path = parse_session_metadata(session_path) - - if not isinstance(training_config_path, Path): - training_config_path = Path(training_config_path) - self._training_config_path: Final[Union[Path, str]] = training_config_path - self._training_config: Final[TrainingConfig] = training_config.load( - self._training_config_path, legacy_training_config - ) - - if not isinstance(lay_down_config_path, Path): - lay_down_config_path = Path(lay_down_config_path) - self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path - self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config) # noqa - - def setup(self) -> None: - """Performs the session setup.""" - if self._training_config.agent_framework == AgentFramework.CUSTOM: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.CUSTOM}") - if self._training_config.agent_identifier == AgentIdentifier.HARDCODED: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Identifier =" f" {AgentIdentifier.HARDCODED}") - if self._training_config.action_type == ActionType.NODE: - # Deterministic Hardcoded Agent with Node Action Space - self._agent_session = HardCodedNodeAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - - elif self._training_config.action_type == ActionType.ACL: - # Deterministic Hardcoded Agent with ACL Action Space - self._agent_session = HardCodedACLAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - - elif self._training_config.action_type == ActionType.ANY: - # Deterministic Hardcoded Agent with ANY Action Space - raise NotImplementedError - - else: - # Invalid AgentIdentifier ActionType combo - raise ValueError - - elif self._training_config.agent_identifier == AgentIdentifier.DO_NOTHING: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Identifier =" f" {AgentIdentifier.DO_NOTHING}") - if self._training_config.action_type == ActionType.NODE: - self._agent_session = DoNothingNodeAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - - elif self._training_config.action_type == ActionType.ACL: - # Deterministic Hardcoded Agent with ACL Action Space - self._agent_session = DoNothingACLAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - - elif self._training_config.action_type == ActionType.ANY: - # Deterministic Hardcoded Agent with ANY Action Space - raise NotImplementedError - - else: - # Invalid AgentIdentifier ActionType combo - raise ValueError - - elif self._training_config.agent_identifier == AgentIdentifier.RANDOM: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Identifier =" f" {AgentIdentifier.RANDOM}") - self._agent_session = RandomAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - elif self._training_config.agent_identifier == AgentIdentifier.DUMMY: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Identifier =" f" {AgentIdentifier.DUMMY}") - self._agent_session = DummyAgent( - self._training_config_path, self._lay_down_config_path, self.session_path - ) - - else: - # Invalid AgentFramework AgentIdentifier combo - raise ValueError - - elif self._training_config.agent_framework == AgentFramework.SB3: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.SB3}") - # Stable Baselines3 Agent - self._agent_session = SB3Agent( - self._training_config_path, - self._lay_down_config_path, - self.session_path, - self.legacy_training_config, - self.legacy_lay_down_config, - ) - - # elif self._training_config.agent_framework == AgentFramework.RLLIB: - # _LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.RLLIB}") - # # Ray RLlib Agent - # self._agent_session = RLlibAgent( - # self._training_config_path, self._lay_down_config_path, self.session_path - # ) - - else: - # Invalid AgentFramework - raise ValueError - - self.session_path: Path = self._agent_session.session_path - self.timestamp_str: str = self._agent_session.timestamp_str - self.learning_path: Path = self._agent_session.learning_path - self.evaluation_path: Path = self._agent_session.evaluation_path - - def learn( - self, - **kwargs: Any, - ) -> None: - """ - Train the agent. - - :param kwargs: Any agent-framework specific key word args. - """ - if not self._training_config.session_type == SessionType.EVAL: - self._agent_session.learn(**kwargs) - - def evaluate( - self, - **kwargs: Any, - ) -> None: - """ - Evaluate the agent. - - :param kwargs: Any agent-framework specific key word args. - """ - if not self._training_config.session_type == SessionType.TRAIN: - self._agent_session.evaluate(**kwargs) - - def close(self) -> None: - """Closes the agent.""" - self._agent_session.close() - - def learn_av_reward_per_episode_dict(self) -> Dict[int, float]: - """Get the learn av reward per episode from file.""" - csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv" - return av_rewards_dict(self.learning_path / csv_file) - - def eval_av_reward_per_episode_dict(self) -> Dict[int, float]: - """Get the eval av reward per episode from file.""" - csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv" - return av_rewards_dict(self.evaluation_path / csv_file) - - def learn_all_transactions_dict(self) -> Dict[Tuple[int, int], Dict[str, Any]]: - """Get the learn all transactions from file.""" - csv_file = f"all_transactions_{self.timestamp_str}.csv" - return all_transactions_dict(self.learning_path / csv_file) - - def eval_all_transactions_dict(self) -> Dict[Tuple[int, int], Dict[str, Any]]: - """Get the eval all transactions from file.""" - csv_file = f"all_transactions_{self.timestamp_str}.csv" - return all_transactions_dict(self.evaluation_path / csv_file) - - def metadata_file_as_dict(self) -> Dict[str, Any]: - """Read the session_metadata.json file and return as a dict.""" - with open(self.session_path / "session_metadata.json", "r") as file: - return json.load(file) diff --git a/src/primaite/setup/old_installation_clean_up.py b/src/primaite/setup/old_installation_clean_up.py deleted file mode 100644 index 412aed60..00000000 --- a/src/primaite/setup/old_installation_clean_up.py +++ /dev/null @@ -1,14 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -from primaite import getLogger - -_LOGGER = getLogger(__name__) - - -def run() -> None: - """Perform the full clean-up.""" - pass - - -if __name__ == "__main__": - run() diff --git a/src/primaite/transactions/__init__.py b/src/primaite/transactions/__init__.py deleted file mode 100644 index 505c5080..00000000 --- a/src/primaite/transactions/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Record data of the system's state and agent's observations and actions.""" diff --git a/src/primaite/transactions/transaction.py b/src/primaite/transactions/transaction.py deleted file mode 100644 index 6b973ca3..00000000 --- a/src/primaite/transactions/transaction.py +++ /dev/null @@ -1,102 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""The Transaction class.""" -from datetime import datetime -from typing import List, Optional, Tuple, TYPE_CHECKING, Union - -from primaite.common.enums import AgentIdentifier - -if TYPE_CHECKING: - import numpy as np - from gymnasium import spaces - - -class Transaction(object): - """Transaction class.""" - - def __init__(self, agent_identifier: AgentIdentifier, episode_number: int, step_number: int) -> None: - """ - Transaction constructor. - - :param agent_identifier: An identifier for the agent in use - :param episode_number: The episode number - :param step_number: The step number - """ - self.timestamp: datetime = datetime.now() - "The datetime of the transaction" - self.agent_identifier: AgentIdentifier = agent_identifier - "The agent identifier" - self.episode_number: int = episode_number - "The episode number" - self.step_number: int = step_number - "The step number" - self.obs_space: "spaces.Space" = None - "The observation space (pre)" - self.obs_space_pre: Optional[Union["np.ndarray", Tuple["np.ndarray"]]] = None - "The observation space before any actions are taken" - self.obs_space_post: Optional[Union["np.ndarray", Tuple["np.ndarray"]]] = None - "The observation space after any actions are taken" - self.reward: Optional[float] = None - "The reward value" - self.action_space: Optional[int] = None - "The action space invoked by the agent" - self.obs_space_description: Optional[List[str]] = None - "The env observation space description" - - def as_csv_data(self) -> Tuple[List, List]: - """ - Converts the Transaction to a csv data row and provides a header. - - :return: A tuple consisting of (header, data). - """ - if isinstance(self.action_space, int): - action_length = self.action_space - else: - action_length = self.action_space.size - - # Create the action space headers array - action_header = [] - for x in range(action_length): - action_header.append("AS_" + str(x)) - - # Open up a csv file - header = ["Timestamp", "Episode", "Step", "Reward"] - header = header + action_header + self.obs_space_description - - row = [ - str(self.timestamp), - str(self.episode_number), - str(self.step_number), - str(self.reward), - ] - row = row + _turn_action_space_to_array(self.action_space) + self.obs_space.tolist() - return header, row - - -def _turn_action_space_to_array(action_space: Union[int, List[int]]) -> List[str]: - """ - Turns action space into a string array so it can be saved to csv. - - :param action_space: The action space - :return: The action space as an array of strings - """ - if isinstance(action_space, list): - return [str(i) for i in action_space] - else: - return [str(action_space)] - - -def _turn_obs_space_to_array(obs_space: "np.ndarray", obs_assets: int, obs_features: int) -> List[str]: - """ - Turns observation space into a string array so it can be saved to csv. - - :param obs_space: The observation space - :param obs_assets: The number of assets (i.e. nodes or links) in the observation space - :param obs_features: The number of features associated with the asset - :return: The observation space as an array of strings - """ - return_array = [] - for x in range(obs_assets): - for y in range(obs_features): - return_array.append(str(obs_space[x][y])) - - return return_array From e8e14ae68a11722a9c016881c71ff6eeb38ad016 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 13:56:02 +0100 Subject: [PATCH 35/53] Comment out tests for primaite v2 --- src/primaite/main.py | 16 +- .../legacy_config_1_DDOS_BASIC.yaml | 170 ------ .../legacy_config_2_DDOS_BASIC.yaml | 362 ------------ .../legacy_config_3_DOS_VERY_BASIC.yaml | 166 ------ .../legacy_config_5_DATA_MANIPULATION.yaml | 534 ------------------ .../legacy_training_config.yaml | 92 --- .../new_training_config.yaml | 114 ---- tests/config/obs_tests/laydown.yaml | 103 ---- tests/config/obs_tests/laydown_ACL.yaml | 86 --- .../main_config_ACCESS_CONTROL_LIST.yaml | 106 ---- .../main_config_LINK_TRAFFIC_LEVELS.yaml | 121 ---- .../main_config_NODE_LINK_TABLE.yaml | 118 ---- .../obs_tests/main_config_NODE_STATUSES.yaml | 115 ---- .../obs_tests/main_config_without_obs.yaml | 108 ---- ...ne_node_states_on_off_lay_down_config.yaml | 117 ---- .../one_node_states_on_off_main_config.yaml | 166 ------ .../ppo_not_seeded_training_config.yaml | 162 ------ tests/config/ppo_seeded_training_config.yaml | 161 ------ .../training_config_main_rllib.yaml | 164 ------ .../training_config_main_sb3.yaml | 164 ------ ..._space_fixed_blue_actions_main_config.yaml | 117 ---- .../single_action_space_lay_down_config.yaml | 45 -- .../single_action_space_main_config.yaml | 116 ---- tests/config/test_random_red_main_config.yaml | 164 ------ tests/config/train_episode_step.yaml | 154 ----- tests/conftest.py | 8 +- tests/test_acl.py | 14 +- tests/test_active_node.py | 8 +- tests/test_full_legacy_config_session.py | 1 + tests/test_lay_down_config.py | 15 +- tests/test_observation_space.py | 8 +- tests/test_primaite_session.py | 4 +- tests/test_red_random_agent_behaviour.py | 5 +- tests/test_resetting_node.py | 13 +- tests/test_reward.py | 1 + .../test_seeding_and_deterministic_session.py | 4 +- tests/test_service_node.py | 8 +- tests/test_session_loading.py | 13 +- tests/test_single_action_space.py | 9 +- tests/test_train_eval_episode_steps.py | 4 +- tests/test_training_config.py | 6 +- 41 files changed, 95 insertions(+), 3767 deletions(-) delete mode 100644 tests/config/legacy_conversion/legacy_config_1_DDOS_BASIC.yaml delete mode 100644 tests/config/legacy_conversion/legacy_config_2_DDOS_BASIC.yaml delete mode 100644 tests/config/legacy_conversion/legacy_config_3_DOS_VERY_BASIC.yaml delete mode 100644 tests/config/legacy_conversion/legacy_config_5_DATA_MANIPULATION.yaml delete mode 100644 tests/config/legacy_conversion/legacy_training_config.yaml delete mode 100644 tests/config/legacy_conversion/new_training_config.yaml delete mode 100644 tests/config/obs_tests/laydown.yaml delete mode 100644 tests/config/obs_tests/laydown_ACL.yaml delete mode 100644 tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml delete mode 100644 tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml delete mode 100644 tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml delete mode 100644 tests/config/obs_tests/main_config_NODE_STATUSES.yaml delete mode 100644 tests/config/obs_tests/main_config_without_obs.yaml delete mode 100644 tests/config/one_node_states_on_off_lay_down_config.yaml delete mode 100644 tests/config/one_node_states_on_off_main_config.yaml delete mode 100644 tests/config/ppo_not_seeded_training_config.yaml delete mode 100644 tests/config/ppo_seeded_training_config.yaml delete mode 100644 tests/config/session_test/training_config_main_rllib.yaml delete mode 100644 tests/config/session_test/training_config_main_sb3.yaml delete mode 100644 tests/config/single_action_space_fixed_blue_actions_main_config.yaml delete mode 100644 tests/config/single_action_space_lay_down_config.yaml delete mode 100644 tests/config/single_action_space_main_config.yaml delete mode 100644 tests/config/test_random_red_main_config.yaml delete mode 100644 tests/config/train_episode_step.yaml diff --git a/src/primaite/main.py b/src/primaite/main.py index 45cd0d8d..0cbcff0e 100644 --- a/src/primaite/main.py +++ b/src/primaite/main.py @@ -5,7 +5,8 @@ from pathlib import Path from typing import Optional, Union from primaite import getLogger -from primaite.primaite_session import PrimaiteSession + +# from primaite.primaite_session import PrimaiteSession _LOGGER = getLogger(__name__) @@ -31,13 +32,14 @@ def run( :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, otherwise False. """ - session = PrimaiteSession( - training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config - ) + # session = PrimaiteSession( + # training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config + # ) - session.setup() - session.learn() - session.evaluate() + # session.setup() + # session.learn() + # session.evaluate() + return NotImplemented if __name__ == "__main__": diff --git a/tests/config/legacy_conversion/legacy_config_1_DDOS_BASIC.yaml b/tests/config/legacy_conversion/legacy_config_1_DDOS_BASIC.yaml deleted file mode 100644 index 5db0ff24..00000000 --- a/tests/config/legacy_conversion/legacy_config_1_DDOS_BASIC.yaml +++ /dev/null @@ -1,170 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- itemType: ACTIONS - type: NODE -- itemType: STEPS - steps: 128 -- itemType: PORTS - portsList: - - port: '80' -- itemType: SERVICES - serviceList: - - name: TCP -- itemType: NODE - id: '1' - name: PC1 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '2' - name: SERVER - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '3' - name: PC2 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '4' - name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.5 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '5' - name: SWITCH2 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.6 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '6' - name: SWITCH3 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.7 - softwareState: GOOD - fileSystemState: GOOD -- itemType: LINK - id: '7' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '4' -- itemType: LINK - id: '8' - name: link2 - bandwidth: 1000000000 - source: '4' - destination: '2' -- itemType: LINK - id: '9' - name: link3 - bandwidth: 1000000000 - source: '2' - destination: '5' -- itemType: LINK - id: '10' - name: link4 - bandwidth: 1000000000 - source: '2' - destination: '6' -- itemType: LINK - id: '11' - name: link5 - bandwidth: 1000000000 - source: '5' - destination: '3' -- itemType: LINK - id: '12' - name: link6 - bandwidth: 1000000000 - source: '6' - destination: '3' -- itemType: GREEN_IER - id: '13' - startStep: 1 - endStep: 128 - load: 100000 - protocol: TCP - port: '80' - source: '3' - destination: '2' - missionCriticality: 5 -- itemType: RED_POL - id: '14' - startStep: 50 - endStep: 50 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_IER - id: '15' - startStep: 60 - endStep: 100 - load: 1000000 - protocol: TCP - port: '80' - source: '1' - destination: '2' - missionCriticality: 0 -- itemType: RED_POL - id: '16' - startStep: 80 - endStep: 80 - targetNodeId: '2' - initiator: IER - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: ACL_RULE - id: '17' - permission: ALLOW - source: ANY - destination: ANY - protocol: ANY - port: ANY diff --git a/tests/config/legacy_conversion/legacy_config_2_DDOS_BASIC.yaml b/tests/config/legacy_conversion/legacy_config_2_DDOS_BASIC.yaml deleted file mode 100644 index 2e791bb1..00000000 --- a/tests/config/legacy_conversion/legacy_config_2_DDOS_BASIC.yaml +++ /dev/null @@ -1,362 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- itemType: ACTIONS - type: NODE -- itemType: STEPS - steps: 128 -- itemType: PORTS - portsList: - - port: '80' -- itemType: SERVICES - serviceList: - - name: TCP -- itemType: NODE - id: '1' - name: PC1 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.11 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '2' - name: PC2 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.12 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '3' - name: PC3 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.13 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '4' - name: PC4 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.20.14 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '5' - name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '6' - name: IDS - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '7' - name: SWITCH2 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '8' - name: LOP1 - baseType: SERVICE - nodeType: LOP - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.12 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '9' - name: SERVER1 - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.14 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '10' - name: SERVER2 - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.20.15 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: LINK - id: '11' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '5' -- itemType: LINK - id: '12' - name: link2 - bandwidth: 1000000000 - source: '2' - destination: '5' -- itemType: LINK - id: '13' - name: link3 - bandwidth: 1000000000 - source: '3' - destination: '5' -- itemType: LINK - id: '14' - name: link4 - bandwidth: 1000000000 - source: '4' - destination: '5' -- itemType: LINK - id: '15' - name: link5 - bandwidth: 1000000000 - source: '5' - destination: '6' -- itemType: LINK - id: '16' - name: link6 - bandwidth: 1000000000 - source: '5' - destination: '8' -- itemType: LINK - id: '17' - name: link7 - bandwidth: 1000000000 - source: '6' - destination: '7' -- itemType: LINK - id: '18' - name: link8 - bandwidth: 1000000000 - source: '8' - destination: '7' -- itemType: LINK - id: '19' - name: link9 - bandwidth: 1000000000 - source: '7' - destination: '9' -- itemType: LINK - id: '20' - name: link10 - bandwidth: 1000000000 - source: '7' - destination: '10' -- itemType: GREEN_IER - id: '21' - startStep: 1 - endStep: 128 - load: 100000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - missionCriticality: 2 -- itemType: GREEN_IER - id: '22' - startStep: 1 - endStep: 128 - load: 100000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - missionCriticality: 2 -- itemType: GREEN_IER - id: '23' - startStep: 1 - endStep: 128 - load: 100000 - protocol: TCP - port: '80' - source: '9' - destination: '3' - missionCriticality: 5 -- itemType: GREEN_IER - id: '24' - startStep: 1 - endStep: 128 - load: 100000 - protocol: TCP - port: '80' - source: '4' - destination: '10' - missionCriticality: 2 -- itemType: ACL_RULE - id: '25' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.10.14 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '26' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.10.14 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '27' - permission: ALLOW - source: 192.168.10.13 - destination: 192.168.10.14 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '28' - permission: ALLOW - source: 192.168.20.14 - destination: 192.168.20.15 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '29' - permission: ALLOW - source: 192.168.10.14 - destination: 192.168.10.13 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '30' - permission: DENY - source: 192.168.10.11 - destination: 192.168.20.15 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '31' - permission: DENY - source: 192.168.10.12 - destination: 192.168.20.15 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '32' - permission: DENY - source: 192.168.10.13 - destination: 192.168.20.15 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '33' - permission: DENY - source: 192.168.20.14 - destination: 192.168.10.14 - protocol: TCP - port: 80 -- itemType: RED_POL - id: '34' - startStep: 20 - endStep: 20 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_POL - id: '35' - startStep: 20 - endStep: 20 - targetNodeId: '2' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_IER - id: '36' - startStep: 30 - endStep: 128 - load: 440000000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - missionCriticality: 0 -- itemType: RED_IER - id: '37' - startStep: 30 - endStep: 128 - load: 440000000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - missionCriticality: 0 -- itemType: RED_POL - id: '38' - startStep: 30 - endStep: 30 - targetNodeId: '9' - initiator: IER - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA diff --git a/tests/config/legacy_conversion/legacy_config_3_DOS_VERY_BASIC.yaml b/tests/config/legacy_conversion/legacy_config_3_DOS_VERY_BASIC.yaml deleted file mode 100644 index 232dd8c7..00000000 --- a/tests/config/legacy_conversion/legacy_config_3_DOS_VERY_BASIC.yaml +++ /dev/null @@ -1,166 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- itemType: ACTIONS - type: NODE -- itemType: STEPS - steps: 256 -- itemType: PORTS - portsList: - - port: '80' -- itemType: SERVICES - serviceList: - - name: TCP -- itemType: NODE - id: '1' - name: PC1 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '2' - name: PC2 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '3' - name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.1 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '4' - name: SERVER1 - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: LINK - id: '5' - name: link1 - bandwidth: 1000000000 - source: '1' - destination: '3' -- itemType: LINK - id: '6' - name: link2 - bandwidth: 1000000000 - source: '2' - destination: '3' -- itemType: LINK - id: '7' - name: link3 - bandwidth: 1000000000 - source: '3' - destination: '4' -- itemType: GREEN_IER - id: '8' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '1' - destination: '4' - missionCriticality: 1 -- itemType: GREEN_IER - id: '9' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '2' - destination: '4' - missionCriticality: 1 -- itemType: GREEN_IER - id: '10' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '4' - destination: '2' - missionCriticality: 5 -- itemType: ACL_RULE - id: '11' - permission: ALLOW - source: 192.168.1.2 - destination: 192.168.1.4 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '12' - permission: ALLOW - source: 192.168.1.3 - destination: 192.168.1.4 - protocol: TCP - port: 80 -- itemType: ACL_RULE - id: '13' - permission: ALLOW - source: 192.168.1.4 - destination: 192.168.1.3 - protocol: TCP - port: 80 -- itemType: RED_POL - id: '14' - startStep: 20 - endStep: 20 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: TCP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_IER - id: '15' - startStep: 30 - endStep: 256 - load: 10000000 - protocol: TCP - port: '80' - source: '1' - destination: '4' - missionCriticality: 0 -- itemType: RED_POL - id: '16' - startStep: 40 - endStep: 40 - targetNodeId: '4' - initiator: IER - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA diff --git a/tests/config/legacy_conversion/legacy_config_5_DATA_MANIPULATION.yaml b/tests/config/legacy_conversion/legacy_config_5_DATA_MANIPULATION.yaml deleted file mode 100644 index 6aa6a4ef..00000000 --- a/tests/config/legacy_conversion/legacy_config_5_DATA_MANIPULATION.yaml +++ /dev/null @@ -1,534 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- itemType: ACTIONS - type: NODE -- itemType: STEPS - steps: 256 -- itemType: PORTS - portsList: - - port: '80' - - port: '1433' - - port: '53' -- itemType: SERVICES - serviceList: - - name: TCP - - name: TCP_SQL - - name: UDP -- itemType: NODE - id: '1' - name: CLIENT_1 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.11 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- itemType: NODE - id: '2' - name: CLIENT_2 - baseType: SERVICE - nodeType: COMPUTER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.12 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: NODE - id: '3' - name: SWITCH_1 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.10.1 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '4' - name: SECURITY_SUITE - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.10 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- itemType: NODE - id: '5' - name: MANAGEMENT_CONSOLE - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.12 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- itemType: NODE - id: '6' - name: SWITCH_2 - baseType: ACTIVE - nodeType: SWITCH - priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.2.1 - softwareState: GOOD - fileSystemState: GOOD -- itemType: NODE - id: '7' - name: WEB_SERVER - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.10 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: TCP_SQL - port: '1433' - state: GOOD -- itemType: NODE - id: '8' - name: DATABASE_SERVER - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.14 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: TCP_SQL - port: '1433' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- itemType: NODE - id: '9' - name: BACKUP_SERVER - baseType: SERVICE - nodeType: SERVER - priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.16 - softwareState: GOOD - fileSystemState: GOOD - services: - - name: TCP - port: '80' - state: GOOD -- itemType: LINK - id: '10' - name: LINK_1 - bandwidth: 1000000000 - source: '1' - destination: '3' -- itemType: LINK - id: '11' - name: LINK_2 - bandwidth: 1000000000 - source: '2' - destination: '3' -- itemType: LINK - id: '12' - name: LINK_3 - bandwidth: 1000000000 - source: '3' - destination: '4' -- itemType: LINK - id: '13' - name: LINK_4 - bandwidth: 1000000000 - source: '3' - destination: '5' -- itemType: LINK - id: '14' - name: LINK_5 - bandwidth: 1000000000 - source: '4' - destination: '6' -- itemType: LINK - id: '15' - name: LINK_6 - bandwidth: 1000000000 - source: '5' - destination: '6' -- itemType: LINK - id: '16' - name: LINK_7 - bandwidth: 1000000000 - source: '6' - destination: '7' -- itemType: LINK - id: '17' - name: LINK_8 - bandwidth: 1000000000 - source: '6' - destination: '8' -- itemType: LINK - id: '18' - name: LINK_9 - bandwidth: 1000000000 - source: '6' - destination: '9' -- itemType: GREEN_IER - id: '19' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '1' - destination: '7' - missionCriticality: 5 -- itemType: GREEN_IER - id: '20' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '7' - destination: '1' - missionCriticality: 5 -- itemType: GREEN_IER - id: '21' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '2' - destination: '7' - missionCriticality: 5 -- itemType: GREEN_IER - id: '22' - startStep: 1 - endStep: 256 - load: 10000 - protocol: TCP - port: '80' - source: '7' - destination: '2' - missionCriticality: 5 -- itemType: GREEN_IER - id: '23' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP_SQL - port: '1433' - source: '7' - destination: '8' - missionCriticality: 5 -- itemType: GREEN_IER - id: '24' - startStep: 1 - endStep: 256 - load: 100000 - protocol: TCP_SQL - port: '1433' - source: '8' - destination: '7' - missionCriticality: 5 -- itemType: GREEN_IER - id: '25' - startStep: 1 - endStep: 256 - load: 50000 - protocol: TCP - port: '80' - source: '1' - destination: '9' - missionCriticality: 2 -- itemType: GREEN_IER - id: '26' - startStep: 1 - endStep: 256 - load: 50000 - protocol: TCP - port: '80' - source: '2' - destination: '9' - missionCriticality: 2 -- itemType: GREEN_IER - id: '27' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '7' - missionCriticality: 1 -- itemType: GREEN_IER - id: '28' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '7' - destination: '5' - missionCriticality: 1 -- itemType: GREEN_IER - id: '29' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '8' - missionCriticality: 1 -- itemType: GREEN_IER - id: '30' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '8' - destination: '5' - missionCriticality: 1 -- itemType: GREEN_IER - id: '31' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '5' - destination: '9' - missionCriticality: 1 -- itemType: GREEN_IER - id: '32' - startStep: 1 - endStep: 256 - load: 5000 - protocol: TCP - port: '80' - source: '9' - destination: '5' - missionCriticality: 1 -- itemType: ACL_RULE - id: '33' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.10 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '34' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.14 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '35' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.14 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '36' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.10 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '37' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.10.11 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '38' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.10.12 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '39' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.2.14 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '40' - permission: ALLOW - source: 192.168.2.14 - destination: 192.168.2.10 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '41' - permission: ALLOW - source: 192.168.10.11 - destination: 192.168.2.16 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '42' - permission: ALLOW - source: 192.168.10.12 - destination: 192.168.2.16 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '43' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.10 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '44' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.14 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '45' - permission: ALLOW - source: 192.168.1.12 - destination: 192.168.2.16 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '46' - permission: ALLOW - source: 192.168.2.10 - destination: 192.168.1.12 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '47' - permission: ALLOW - source: 192.168.2.14 - destination: 192.168.1.12 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '48' - permission: ALLOW - source: 192.168.2.16 - destination: 192.168.1.12 - protocol: ANY - port: ANY -- itemType: ACL_RULE - id: '49' - permission: DENY - source: ANY - destination: ANY - protocol: ANY - port: ANY -- itemType: RED_POL - id: '50' - startStep: 50 - endStep: 50 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: UDP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_IER - id: '51' - startStep: 75 - endStep: 105 - load: 10000 - protocol: UDP - port: '53' - source: '1' - destination: '8' - missionCriticality: 0 -- itemType: RED_POL - id: '52' - startStep: 100 - endStep: 100 - targetNodeId: '8' - initiator: IER - type: SERVICE - protocol: UDP - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_POL - id: '53' - startStep: 105 - endStep: 105 - targetNodeId: '8' - initiator: SERVICE - type: FILE - protocol: NA - state: CORRUPT - sourceNodeId: '8' - sourceNodeService: UDP - sourceNodeServiceState: COMPROMISED -- itemType: RED_POL - id: '54' - startStep: 105 - endStep: 105 - targetNodeId: '8' - initiator: SERVICE - type: SERVICE - protocol: TCP_SQL - state: COMPROMISED - sourceNodeId: '8' - sourceNodeService: UDP - sourceNodeServiceState: COMPROMISED -- itemType: RED_POL - id: '55' - startStep: 125 - endStep: 125 - targetNodeId: '7' - initiator: SERVICE - type: SERVICE - protocol: TCP - state: OVERWHELMED - sourceNodeId: '8' - sourceNodeService: TCP_SQL - sourceNodeServiceState: COMPROMISED diff --git a/tests/config/legacy_conversion/legacy_training_config.yaml b/tests/config/legacy_conversion/legacy_training_config.yaml deleted file mode 100644 index 3477e6e0..00000000 --- a/tests/config/legacy_conversion/legacy_training_config.yaml +++ /dev/null @@ -1,92 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Main Config File - -# Generic config values -# Choose one of these (dependent on Agent being trained) -# "STABLE_BASELINES3_PPO" -# "STABLE_BASELINES3_A2C" -# "GENERIC" -agentIdentifier: STABLE_BASELINES3_A2C -# Number of episodes to run per session -numEpisodes: 10 -# Time delay between steps (for generic agents) -timeDelay: 10 -# Filename of the scenario / laydown -configFilename: config_5_DATA_MANIPULATION.yaml -# Type of session to be run (TRAINING or EVALUATION) -sessionType: TRAINING -# Determine whether to load an agent from file -loadAgent: False -# File path and file name of agent if you're loading one in -agentLoadFile: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observationSpaceHighValue: 1000000000 - -# Reward values -# Generic -allOk: 0 -# Node Hardware State -offShouldBeOn: -10 -offShouldBeResetting: -5 -onShouldBeOff: -2 -onShouldBeResetting: -5 -resettingShouldBeOn: -5 -resettingShouldBeOff: -2 -resetting: -3 -# Node Software or Service State -goodShouldBePatching: 2 -goodShouldBeCompromised: 5 -goodShouldBeOverwhelmed: 5 -patchingShouldBeGood: -5 -patchingShouldBeCompromised: 2 -patchingShouldBeOverwhelmed: 2 -patching: -3 -compromisedShouldBeGood: -20 -compromisedShouldBePatching: -20 -compromisedShouldBeOverwhelmed: -20 -compromised: -20 -overwhelmedShouldBeGood: -20 -overwhelmedShouldBePatching: -20 -overwhelmedShouldBeCompromised: -20 -overwhelmed: -20 -# Node File System State -goodShouldBeRepairing: 2 -goodShouldBeRestoring: 2 -goodShouldBeCorrupt: 5 -goodShouldBeDestroyed: 10 -repairingShouldBeGood: -5 -repairingShouldBeRestoring: 2 -repairingShouldBeCorrupt: 2 -repairingShouldBeDestroyed: 0 -repairing: -3 -restoringShouldBeGood: -10 -restoringShouldBeRepairing: -2 -restoringShouldBeCorrupt: 1 -restoringShouldBeDestroyed: 2 -restoring: -6 -corruptShouldBeGood: -10 -corruptShouldBeRepairing: -10 -corruptShouldBeRestoring: -10 -corruptShouldBeDestroyed: 2 -corrupt: -10 -destroyedShouldBeGood: -20 -destroyedShouldBeRepairing: -20 -destroyedShouldBeRestoring: -20 -destroyedShouldBeCorrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -redIerRunning: -5 -greenIerBlocked: -10 - -# Patching / Reset durations -osPatchingDuration: 5 # The time taken to patch the OS -nodeResetDuration: 5 # The time taken to reset a node (hardware) -nodeBootingDuration: 3 # The Time taken to turn on the node -nodeShutdownDuration: 2 # The time taken to turn off the node -servicePatchingDuration: 5 # The time taken to patch a service -fileSystemRepairingLimit: 5 # The time take to repair the file system -fileSystemRestoringLimit: 5 # The time take to restore the file system -fileSystemScanningLimit: 5 # The time taken to scan the file system diff --git a/tests/config/legacy_conversion/new_training_config.yaml b/tests/config/legacy_conversion/new_training_config.yaml deleted file mode 100644 index 1991eb06..00000000 --- a/tests/config/legacy_conversion/new_training_config.yaml +++ /dev/null @@ -1,114 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Main Config File - -# Generic config values - -# Sets which agent algorithm framework will be used: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray[RLlib]) -# "NONE" (Custom Agent) -agent_framework: SB3 - -# Sets which Red Agent algo/class will be used: -# "PPO" (Proximal Policy Optimization) -# "A2C" (Advantage Actor Critic) -# "HARDCODED" (Custom Agent) -# "RANDOM" (Random Action) -agent_identifier: PPO - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 1 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 - - -# Time delay between steps (for generic agents) -time_delay: 10 -# Type of session to be run (TRAINING or EVALUATION) -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/obs_tests/laydown.yaml b/tests/config/obs_tests/laydown.yaml deleted file mode 100644 index e358d0d2..00000000 --- a/tests/config/obs_tests/laydown.yaml +++ /dev/null @@ -1,103 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- item_type: PORTS - ports_list: - - port: '80' - - port: '53' -- item_type: SERVICES - service_list: - - name: TCP - - name: UDP - -######################################## -# Nodes -- item_type: NODE - node_id: '1' - name: PC1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.1 - software_state: COMPROMISED - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: GOOD -- item_type: NODE - node_id: '2' - name: SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.2 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: UDP - port: '53' - state: OVERWHELMED -- item_type: NODE - node_id: '3' - name: SWITCH1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.3 - software_state: GOOD - file_system_state: GOOD - -######################################## -# Links -- item_type: LINK - id: '4' - name: link1 - bandwidth: 1000 - source: '1' - destination: '3' -- item_type: LINK - id: '5' - name: link2 - bandwidth: 1000 - source: '3' - destination: '2' - -######################################### -# IERS -- item_type: GREEN_IER - id: '5' - start_step: 0 - end_step: 5 - load: 999 - protocol: TCP - port: '80' - source: '1' - destination: '2' - mission_criticality: 5 - -######################################### -# ACL Rules -- item_type: ACL_RULE - id: '6' - permission: ALLOW - source: 192.168.1.1 - destination: 192.168.1.2 - protocol: TCP - port: 80 - position: 0 -- item_type: ACL_RULE - id: '7' - permission: ALLOW - source: 192.168.1.2 - destination: 192.168.1.1 - protocol: TCP - port: 80 - position: 0 diff --git a/tests/config/obs_tests/laydown_ACL.yaml b/tests/config/obs_tests/laydown_ACL.yaml deleted file mode 100644 index cffd8b1c..00000000 --- a/tests/config/obs_tests/laydown_ACL.yaml +++ /dev/null @@ -1,86 +0,0 @@ -- item_type: PORTS - ports_list: - - port: '80' - - port: '21' -- item_type: SERVICES - service_list: - - name: TCP - - name: FTP - -######################################## -# Nodes -- item_type: NODE - node_id: '1' - name: PC1 - node_class: SERVICE - node_type: COMPUTER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.1 - software_state: COMPROMISED - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: FTP - port: '21' - state: GOOD -- item_type: NODE - node_id: '2' - name: SERVER - node_class: SERVICE - node_type: SERVER - priority: P5 - hardware_state: 'ON' - ip_address: 192.168.1.2 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: GOOD - - name: FTP - port: '21' - state: OVERWHELMED -- item_type: NODE - node_id: '3' - name: SWITCH1 - node_class: ACTIVE - node_type: SWITCH - priority: P2 - hardware_state: 'ON' - ip_address: 192.168.1.3 - software_state: GOOD - file_system_state: GOOD - -######################################## -# Links -- item_type: LINK - id: '4' - name: link1 - bandwidth: 1000 - source: '1' - destination: '3' -- item_type: LINK - id: '5' - name: link2 - bandwidth: 1000 - source: '3' - destination: '2' - -######################################### -# IERS -- item_type: GREEN_IER - id: '5' - start_step: 0 - end_step: 5 - load: 999 - protocol: TCP - port: '80' - source: '1' - destination: '2' - mission_criticality: 5 - -######################################### -# ACL Rules diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml deleted file mode 100644 index 927c9f44..00000000 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ /dev/null @@ -1,106 +0,0 @@ -# Main Config File - -# Generic config values -# Choose one of these (dependent on Agent being trained) -# "STABLE_BASELINES3_PPO" -# "STABLE_BASELINES3_A2C" -# "GENERIC" -agent_framework: SB3 -agent_identifier: PPO -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 -# Number of time_steps for training per episode -num_train_steps: 5 - -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) -implicit_acl_rule: DENY -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 3 - -observation_space: - components: - - name: ACCESS_CONTROL_LIST - -# Time delay between steps (for generic agents) -time_delay: 1 - -# Type of session to be run (TRAINING or EVALUATION) -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1_000_000_000 - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml b/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml deleted file mode 100644 index 805ab31e..00000000 --- a/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml +++ /dev/null @@ -1,121 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: A2C - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 - -# Number of time_steps for training per episode -num_train_steps: 5 - -observation_space: - components: - - name: LINK_TRAFFIC_LEVELS - options: - combine_service_traffic: false - quantisation_levels: 8 - -# Time delay between steps (for generic agents) -time_delay: 1 - -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) -implicit_acl_rule: ALLOW -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 4 - -# Type of session to be run (TRAINING or EVALUATION) -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1_000_000_000 - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml b/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml deleted file mode 100644 index 535558aa..00000000 --- a/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml +++ /dev/null @@ -1,118 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: RANDOM - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 - -# Number of time_steps for training per episode -num_train_steps: 5 - -observation_space: - components: - - name: NODE_LINK_TABLE - -# Time delay between steps (for generic agents) -time_delay: 1 -# Filename of the scenario / laydown - -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) -implicit_acl_rule: ALLOW -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 4 - -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1_000_000_000 - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/obs_tests/main_config_NODE_STATUSES.yaml b/tests/config/obs_tests/main_config_NODE_STATUSES.yaml deleted file mode 100644 index d1319c35..00000000 --- a/tests/config/obs_tests/main_config_NODE_STATUSES.yaml +++ /dev/null @@ -1,115 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: RANDOM - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 - -# Number of time_steps for training per episode -num_train_steps: 5 - - -observation_space: - components: - - name: NODE_STATUSES - - -# Time delay between steps (for generic agents) -time_delay: 1 - -# Type of session to be run (TRAINING or EVALUATION) -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1_000_000_000 - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/obs_tests/main_config_without_obs.yaml b/tests/config/obs_tests/main_config_without_obs.yaml deleted file mode 100644 index 26457c84..00000000 --- a/tests/config/obs_tests/main_config_without_obs.yaml +++ /dev/null @@ -1,108 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: RANDOM - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 - -# Number of time_steps for training per episode -num_train_steps: 5 -# Time delay between steps (for generic agents) -time_delay: 1 -# Type of session to be run (TRAINING or EVALUATION) -session_type: TRAIN -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1_000_000_000 -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) -implicit_acl_rule: DENY -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/one_node_states_on_off_lay_down_config.yaml b/tests/config/one_node_states_on_off_lay_down_config.yaml deleted file mode 100644 index 0f572d8d..00000000 --- a/tests/config/one_node_states_on_off_lay_down_config.yaml +++ /dev/null @@ -1,117 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- item_type: PORTS - ports_list: - - port: '21' -- item_type: SERVICES - service_list: - - name: ftp -- item_type: NODE - node_id: '1' - name: node - node_class: SERVICE - node_type: COMPUTER - priority: P1 - hardware_state: 'ON' - ip_address: 192.168.0.1 - software_state: GOOD - file_system_state: GOOD - services: - - name: ftp - port: '21' - state: GOOD -- item_type: RED_POL - id: '1' - start_step: 1 - end_step: 3 - targetNodeId: '1' - initiator: DIRECT - type: FILE - protocol: NA - state: CORRUPT - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '2' - start_step: 3 - end_step: 15 - targetNodeId: '1' - initiator: DIRECT - type: FILE - protocol: NA - state: GOOD - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '3' - start_step: 4 - end_step: 6 - targetNodeId: '1' - initiator: DIRECT - type: OPERATING - protocol: NA - state: 'OFF' - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '4' - start_step: 6 - end_step: 15 - targetNodeId: '1' - initiator: DIRECT - type: OPERATING - protocol: NA - state: 'ON' - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '5' - start_step: 7 - end_step: 9 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: ftp - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '6' - start_step: 9 - end_step: 15 - targetNodeId: '1' - initiator: DIRECT - type: SERVICE - protocol: ftp - state: GOOD - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '7' - start_step: 10 - end_step: 12 - targetNodeId: '1' - initiator: DIRECT - type: OS - protocol: NA - state: COMPROMISED - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- item_type: RED_POL - id: '8' - start_step: 12 - end_step: 15 - targetNodeId: '1' - initiator: DIRECT - type: OS - protocol: NA - state: GOOD - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA diff --git a/tests/config/one_node_states_on_off_main_config.yaml b/tests/config/one_node_states_on_off_main_config.yaml deleted file mode 100644 index 10af7a1f..00000000 --- a/tests/config/one_node_states_on_off_main_config.yaml +++ /dev/null @@ -1,166 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: DUMMY - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: null - -# Set whether the agent will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - # flatten: true - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 1 - -# Number of time_steps for evaluation per episode -num_eval_steps: 15 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 10 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -implicit_acl_rule: DENY -max_number_acl_rules: 10 -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/ppo_not_seeded_training_config.yaml b/tests/config/ppo_not_seeded_training_config.yaml deleted file mode 100644 index fac2fe95..00000000 --- a/tests/config/ppo_not_seeded_training_config.yaml +++ /dev/null @@ -1,162 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: None - -# Set whether the agent evaluation will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - # - name: ACCESS_CONTROL_LIST -# Number of episodes to run per session -num_train_episodes: 10 - -# Number of time_steps per episode -num_train_steps: 256 - -# Number of episodes to run per session -num_eval_episodes: 10 - -# Number of time_steps per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 0 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0.0000 -# Node Hardware State -off_should_be_on: -0.001 -off_should_be_resetting: -0.0005 -on_should_be_off: -0.0002 -on_should_be_resetting: -0.0005 -resetting_should_be_on: -0.0005 -resetting_should_be_off: -0.0002 -resetting: -0.0003 -# Node Software or Service State -good_should_be_patching: 0.0002 -good_should_be_compromised: 0.0005 -good_should_be_overwhelmed: 0.0005 -patching_should_be_good: -0.0005 -patching_should_be_compromised: 0.0002 -patching_should_be_overwhelmed: 0.0002 -patching: -0.0003 -compromised_should_be_good: -0.002 -compromised_should_be_patching: -0.002 -compromised_should_be_overwhelmed: -0.002 -compromised: -0.002 -overwhelmed_should_be_good: -0.002 -overwhelmed_should_be_patching: -0.002 -overwhelmed_should_be_compromised: -0.002 -overwhelmed: -0.002 -# Node File System State -good_should_be_repairing: 0.0002 -good_should_be_restoring: 0.0002 -good_should_be_corrupt: 0.0005 -good_should_be_destroyed: 0.001 -repairing_should_be_good: -0.0005 -repairing_should_be_restoring: 0.0002 -repairing_should_be_corrupt: 0.0002 -repairing_should_be_destroyed: 0.0000 -repairing: -0.0003 -restoring_should_be_good: -0.001 -restoring_should_be_repairing: -0.0002 -restoring_should_be_corrupt: 0.0001 -restoring_should_be_destroyed: 0.0002 -restoring: -0.0006 -corrupt_should_be_good: -0.001 -corrupt_should_be_repairing: -0.001 -corrupt_should_be_restoring: -0.001 -corrupt_should_be_destroyed: 0.0002 -corrupt: -0.001 -destroyed_should_be_good: -0.002 -destroyed_should_be_repairing: -0.002 -destroyed_should_be_restoring: -0.002 -destroyed_should_be_corrupt: -0.002 -destroyed: -0.002 -scanning: -0.0002 -# IER status -red_ier_running: -0.0005 -green_ier_blocked: -0.001 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/ppo_seeded_training_config.yaml b/tests/config/ppo_seeded_training_config.yaml deleted file mode 100644 index e4d4fe5b..00000000 --- a/tests/config/ppo_seeded_training_config.yaml +++ /dev/null @@ -1,161 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: 67890 - -# Set whether the agent evaluation will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: True - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS -# Number of episodes to run per session -num_train_episodes: 10 - -# Number of time_steps per episode -num_train_steps: 256 - -# Number of episodes to run per session -num_eval_episodes: 1 - -# Number of time_steps per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 0 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/session_test/training_config_main_rllib.yaml b/tests/config/session_test/training_config_main_rllib.yaml deleted file mode 100644 index 374c6ac5..00000000 --- a/tests/config/session_test/training_config_main_rllib.yaml +++ /dev/null @@ -1,164 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: RLLIB - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: null - -# Set whether the agent will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - # flatten: true - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 3 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 10 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -0.001 -off_should_be_resetting: -0.0005 -on_should_be_off: -0.0002 -on_should_be_resetting: -0.0005 -resetting_should_be_on: -0.0005 -resetting_should_be_off: -0.0002 -resetting: -0.0003 -# Node Software or Service State -good_should_be_patching: 0.0002 -good_should_be_compromised: 0.0005 -good_should_be_overwhelmed: 0.0005 -patching_should_be_good: -0.0005 -patching_should_be_compromised: 0.0002 -patching_should_be_overwhelmed: 0.0002 -patching: -0.0003 -compromised_should_be_good: -0.002 -compromised_should_be_patching: -0.002 -compromised_should_be_overwhelmed: -0.002 -compromised: -0.002 -overwhelmed_should_be_good: -0.002 -overwhelmed_should_be_patching: -0.002 -overwhelmed_should_be_compromised: -0.002 -overwhelmed: -0.002 -# Node File System State -good_should_be_repairing: 0.0002 -good_should_be_restoring: 0.0002 -good_should_be_corrupt: 0.0005 -good_should_be_destroyed: 0.001 -repairing_should_be_good: -0.0005 -repairing_should_be_restoring: 0.0002 -repairing_should_be_corrupt: 0.0002 -repairing_should_be_destroyed: 0.0000 -repairing: -0.0003 -restoring_should_be_good: -0.001 -restoring_should_be_repairing: -0.0002 -restoring_should_be_corrupt: 0.0001 -restoring_should_be_destroyed: 0.0002 -restoring: -0.0006 -corrupt_should_be_good: -0.001 -corrupt_should_be_repairing: -0.001 -corrupt_should_be_restoring: -0.001 -corrupt_should_be_destroyed: 0.0002 -corrupt: -0.001 -destroyed_should_be_good: -0.002 -destroyed_should_be_repairing: -0.002 -destroyed_should_be_restoring: -0.002 -destroyed_should_be_corrupt: -0.002 -destroyed: -0.002 -scanning: -0.0002 -# IER status -red_ier_running: -0.0005 -green_ier_blocked: -0.001 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/session_test/training_config_main_sb3.yaml b/tests/config/session_test/training_config_main_sb3.yaml deleted file mode 100644 index 733105ea..00000000 --- a/tests/config/session_test/training_config_main_sb3.yaml +++ /dev/null @@ -1,164 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: null - -# Set whether the agent will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - # flatten: true - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 3 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 10 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -0.001 -off_should_be_resetting: -0.0005 -on_should_be_off: -0.0002 -on_should_be_resetting: -0.0005 -resetting_should_be_on: -0.0005 -resetting_should_be_off: -0.0002 -resetting: -0.0003 -# Node Software or Service State -good_should_be_patching: 0.0002 -good_should_be_compromised: 0.0005 -good_should_be_overwhelmed: 0.0005 -patching_should_be_good: -0.0005 -patching_should_be_compromised: 0.0002 -patching_should_be_overwhelmed: 0.0002 -patching: -0.0003 -compromised_should_be_good: -0.002 -compromised_should_be_patching: -0.002 -compromised_should_be_overwhelmed: -0.002 -compromised: -0.002 -overwhelmed_should_be_good: -0.002 -overwhelmed_should_be_patching: -0.002 -overwhelmed_should_be_compromised: -0.002 -overwhelmed: -0.002 -# Node File System State -good_should_be_repairing: 0.0002 -good_should_be_restoring: 0.0002 -good_should_be_corrupt: 0.0005 -good_should_be_destroyed: 0.001 -repairing_should_be_good: -0.0005 -repairing_should_be_restoring: 0.0002 -repairing_should_be_corrupt: 0.0002 -repairing_should_be_destroyed: 0.0000 -repairing: -0.0003 -restoring_should_be_good: -0.001 -restoring_should_be_repairing: -0.0002 -restoring_should_be_corrupt: 0.0001 -restoring_should_be_destroyed: 0.0002 -restoring: -0.0006 -corrupt_should_be_good: -0.001 -corrupt_should_be_repairing: -0.001 -corrupt_should_be_restoring: -0.001 -corrupt_should_be_destroyed: 0.0002 -corrupt: -0.001 -destroyed_should_be_good: -0.002 -destroyed_should_be_repairing: -0.002 -destroyed_should_be_restoring: -0.002 -destroyed_should_be_corrupt: -0.002 -destroyed: -0.002 -scanning: -0.0002 -# IER status -red_ier_running: -0.0005 -green_ier_blocked: -0.001 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml deleted file mode 100644 index 6210cf3e..00000000 --- a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml +++ /dev/null @@ -1,117 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: RANDOM - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 1 - -# Number of time_steps for training per episode -num_train_steps: 15 - -# Time delay between steps (for generic agents) -time_delay: 1 -# Type of session to be run (TRAINING or EVALUATION) -session_type: EVAL -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) -implicit_acl_rule: DENY -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 10 - -observation_space: - components: - - name: ACCESS_CONTROL_LIST - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# Reward values -# Generic -all_ok: 0 -# Node Operating State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node O/S or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/single_action_space_lay_down_config.yaml b/tests/config/single_action_space_lay_down_config.yaml deleted file mode 100644 index 9103e2b7..00000000 --- a/tests/config/single_action_space_lay_down_config.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -- item_type: PORTS - ports_list: - - port: '80' -- item_type: SERVICES - service_list: - - name: TCP -- item_type: NODE - node_id: '1' - name: node - node_class: SERVICE - node_type: COMPUTER - priority: P1 - hardware_state: 'ON' - ip_address: 192.168.0.14 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: COMPROMISED -- item_type: NODE - node_id: '2' - name: server_1 - node_class: SERVICE - node_type: SERVER - priority: P1 - hardware_state: 'ON' - ip_address: 192.168.0.1 - software_state: GOOD - file_system_state: GOOD - services: - - name: TCP - port: '80' - state: COMPROMISED -- item_type: RED_IER - id: '3' - start_step: 2 - end_step: 15 - load: 1000 - protocol: TCP - port: CORRUPT - source: '1' - destination: '2' - mission_criticality: 0 diff --git a/tests/config/single_action_space_main_config.yaml b/tests/config/single_action_space_main_config.yaml deleted file mode 100644 index 67eaf49d..00000000 --- a/tests/config/single_action_space_main_config.yaml +++ /dev/null @@ -1,116 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: CUSTOM - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: RANDOM - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: ANY -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 10 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 -# Time delay between steps (for generic agents) -time_delay: 1 -# Type of session to be run (TRAINING or EVALUATION) -session_type: EVAL -# Determine whether to load an agent from file -load_agent: False -# File path and file name of agent if you're loading one in -agent_load_file: C:\[Path]\[agent_saved_filename.zip] - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) -implicit_acl_rule: DENY -max_number_acl_rules: 10 -# Reward values -# Generic -all_ok: 0 -# Node Operating State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node O/S or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/test_random_red_main_config.yaml b/tests/config/test_random_red_main_config.yaml deleted file mode 100644 index 310c9dc6..00000000 --- a/tests/config/test_random_red_main_config.yaml +++ /dev/null @@ -1,164 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: True - -# The (integer) seed to be used in random number generation -# Default is None (null) -seed: null - -# Set whether the agent will be deterministic instead of stochastic -# Options are: -# True -# False -deterministic: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - # flatten: true - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - - -# Number of episodes for training to run per session -num_train_episodes: 10 - -# Number of time_steps for training per episode -num_train_steps: 256 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 1 - -# Number of time_steps for evaluation per episode -num_eval_steps: 256 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 10 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -0.001 -off_should_be_resetting: -0.0005 -on_should_be_off: -0.0002 -on_should_be_resetting: -0.0005 -resetting_should_be_on: -0.0005 -resetting_should_be_off: -0.0002 -resetting: -0.0003 -# Node Software or Service State -good_should_be_patching: 0.0002 -good_should_be_compromised: 0.0005 -good_should_be_overwhelmed: 0.0005 -patching_should_be_good: -0.0005 -patching_should_be_compromised: 0.0002 -patching_should_be_overwhelmed: 0.0002 -patching: -0.0003 -compromised_should_be_good: -0.002 -compromised_should_be_patching: -0.002 -compromised_should_be_overwhelmed: -0.002 -compromised: -0.002 -overwhelmed_should_be_good: -0.002 -overwhelmed_should_be_patching: -0.002 -overwhelmed_should_be_compromised: -0.002 -overwhelmed: -0.002 -# Node File System State -good_should_be_repairing: 0.0002 -good_should_be_restoring: 0.0002 -good_should_be_corrupt: 0.0005 -good_should_be_destroyed: 0.001 -repairing_should_be_good: -0.0005 -repairing_should_be_restoring: 0.0002 -repairing_should_be_corrupt: 0.0002 -repairing_should_be_destroyed: 0.0000 -repairing: -0.0003 -restoring_should_be_good: -0.001 -restoring_should_be_repairing: -0.0002 -restoring_should_be_corrupt: 0.0001 -restoring_should_be_destroyed: 0.0002 -restoring: -0.0006 -corrupt_should_be_good: -0.001 -corrupt_should_be_repairing: -0.001 -corrupt_should_be_restoring: -0.001 -corrupt_should_be_destroyed: 0.0002 -corrupt: -0.001 -destroyed_should_be_good: -0.002 -destroyed_should_be_repairing: -0.002 -destroyed_should_be_restoring: -0.002 -destroyed_should_be_corrupt: -0.002 -destroyed: -0.002 -scanning: -0.0002 -# IER status -red_ier_running: -0.0005 -green_ier_blocked: -0.001 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/config/train_episode_step.yaml b/tests/config/train_episode_step.yaml deleted file mode 100644 index a86e0f62..00000000 --- a/tests/config/train_episode_step.yaml +++ /dev/null @@ -1,154 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -# Training Config File - -# Sets which agent algorithm framework will be used. -# Options are: -# "SB3" (Stable Baselines3) -# "RLLIB" (Ray RLlib) -# "CUSTOM" (Custom Agent) -agent_framework: SB3 - -# Sets which deep learning framework will be used (by RLlib ONLY). -# Default is TF (Tensorflow). -# Options are: -# "TF" (Tensorflow) -# TF2 (Tensorflow 2.X) -# TORCH (PyTorch) -deep_learning_framework: TF2 - -# Sets which Agent class will be used. -# Options are: -# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework) -# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework) -# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type) -# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type) -# "RANDOM" (primaite.agents.simple.RandomAgent) -# "DUMMY" (primaite.agents.simple.DummyAgent) -agent_identifier: PPO - -# Sets whether Red Agent POL and IER is randomised. -# Options are: -# True -# False -random_red_agent: False - -# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC. -# Options are: -# "BASIC" (The current observation space only) -# "FULL" (Full environment view with actions taken and reward feedback) -hard_coded_agent_view: FULL - -# Sets How the Action Space is defined: -# "NODE" -# "ACL" -# "ANY" node and acl actions -action_type: NODE -# observation space -observation_space: - # flatten: true - components: - - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS - - -# Number of episodes for training to run per session -num_train_episodes: 3 - -# Number of time_steps for training per episode -num_train_steps: 25 - -# Number of episodes for evaluation to run per session -num_eval_episodes: 1 - -# Number of time_steps for evaluation per episode -num_eval_steps: 17 - -# Sets how often the agent will save a checkpoint (every n time episodes). -# Set to 0 if no checkpoints are required. Default is 10 -checkpoint_every_n_episodes: 0 - -# Time delay (milliseconds) between steps for CUSTOM agents. -time_delay: 5 - -# Type of session to be run. Options are: -# "TRAIN" (Trains an agent) -# "EVAL" (Evaluates an agent) -# "TRAIN_EVAL" (Trains then evaluates an agent) -session_type: TRAIN_EVAL - -# Environment config values -# The high value for the observation space -observation_space_high_value: 1000000000 - -# The Stable Baselines3 learn/eval output verbosity level: -# Options are: -# "NONE" (No Output) -# "INFO" (Info Messages (such as devices and wrappers used)) -# "DEBUG" (All Messages) -sb3_output_verbose_level: NONE - -# Reward values -# Generic -all_ok: 0 -# Node Hardware State -off_should_be_on: -10 -off_should_be_resetting: -5 -on_should_be_off: -2 -on_should_be_resetting: -5 -resetting_should_be_on: -5 -resetting_should_be_off: -2 -resetting: -3 -# Node Software or Service State -good_should_be_patching: 2 -good_should_be_compromised: 5 -good_should_be_overwhelmed: 5 -patching_should_be_good: -5 -patching_should_be_compromised: 2 -patching_should_be_overwhelmed: 2 -patching: -3 -compromised_should_be_good: -20 -compromised_should_be_patching: -20 -compromised_should_be_overwhelmed: -20 -compromised: -20 -overwhelmed_should_be_good: -20 -overwhelmed_should_be_patching: -20 -overwhelmed_should_be_compromised: -20 -overwhelmed: -20 -# Node File System State -good_should_be_repairing: 2 -good_should_be_restoring: 2 -good_should_be_corrupt: 5 -good_should_be_destroyed: 10 -repairing_should_be_good: -5 -repairing_should_be_restoring: 2 -repairing_should_be_corrupt: 2 -repairing_should_be_destroyed: 0 -repairing: -3 -restoring_should_be_good: -10 -restoring_should_be_repairing: -2 -restoring_should_be_corrupt: 1 -restoring_should_be_destroyed: 2 -restoring: -6 -corrupt_should_be_good: -10 -corrupt_should_be_repairing: -10 -corrupt_should_be_restoring: -10 -corrupt_should_be_destroyed: 2 -corrupt: -10 -destroyed_should_be_good: -20 -destroyed_should_be_repairing: -20 -destroyed_should_be_restoring: -20 -destroyed_should_be_corrupt: -20 -destroyed: -20 -scanning: -2 -# IER status -red_ier_running: -5 -green_ier_blocked: -10 - -# Patching / Reset durations -os_patching_duration: 5 # The time taken to patch the OS -node_reset_duration: 5 # The time taken to reset a node (hardware) -service_patching_duration: 5 # The time taken to patch a service -file_system_repairing_limit: 5 # The time take to repair the file system -file_system_restoring_limit: 5 # The time take to restore the file system -file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/conftest.py b/tests/conftest.py index d8c9cc50..425acc09 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,8 +10,9 @@ from unittest.mock import patch import pytest from primaite import getLogger -from primaite.environment.primaite_env import Primaite -from primaite.primaite_session import PrimaiteSession + +# from primaite.environment.primaite_env import Primaite +# from primaite.primaite_session import PrimaiteSession from primaite.simulator.network.container import Network from primaite.simulator.network.networks import arcd_uc2_network from primaite.simulator.network.transmission.transport_layer import Port @@ -53,6 +54,7 @@ def file_system() -> FileSystem: return Node(hostname="fs_node").file_system +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 # PrimAITE v2 stuff class TempPrimaiteSession(PrimaiteSession): """ @@ -82,6 +84,7 @@ class TempPrimaiteSession(PrimaiteSession): _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.fixture def temp_primaite_session(request): """ @@ -136,6 +139,7 @@ def temp_primaite_session(request): return TempPrimaiteSession(training_config_path, lay_down_config_path) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.fixture def temp_session_path() -> Path: """ diff --git a/tests/test_acl.py b/tests/test_acl.py index d8357cf6..a0dfb997 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -1,11 +1,12 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK """Used to tes the ACL functions.""" -from primaite.acl.access_control_list import AccessControlList -from primaite.acl.acl_rule import ACLRule -from primaite.common.enums import RulePermissionType +# from primaite.acl.access_control_list import AccessControlList +# from primaite.acl.acl_rule import ACLRule +# from primaite.common.enums import RulePermissionType +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_acl_address_match_1(): """Test that matching IP addresses produce True.""" acl = AccessControlList(RulePermissionType.DENY, 10) @@ -15,6 +16,7 @@ def test_acl_address_match_1(): assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_acl_address_match_2(): """Test that mismatching IP addresses produce False.""" acl = AccessControlList(RulePermissionType.DENY, 10) @@ -24,6 +26,7 @@ def test_acl_address_match_2(): assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_acl_address_match_3(): """Test the ANY condition for source IP addresses produce True.""" acl = AccessControlList(RulePermissionType.DENY, 10) @@ -33,6 +36,7 @@ def test_acl_address_match_3(): assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_acl_address_match_4(): """Test the ANY condition for dest IP addresses produce True.""" acl = AccessControlList(RulePermissionType.DENY, 10) @@ -42,6 +46,7 @@ def test_acl_address_match_4(): assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_check_acl_block_affirmative(): """Test the block function (affirmative).""" # Create the Access Control List @@ -66,6 +71,7 @@ def test_check_acl_block_affirmative(): assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_check_acl_block_negative(): """Test the block function (negative).""" # Create the Access Control List @@ -91,6 +97,7 @@ def test_check_acl_block_negative(): assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == True +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_rule_hash(): """Test the rule hash.""" # Create the Access Control List @@ -104,6 +111,7 @@ def test_rule_hash(): assert hash_value_local == hash_value_remote +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_delete_rule(): """Adds 3 rules and deletes 1 rule and checks its deletion.""" # Create the Access Control List diff --git a/tests/test_active_node.py b/tests/test_active_node.py index 44d38313..cf532bb8 100644 --- a/tests/test_active_node.py +++ b/tests/test_active_node.py @@ -2,10 +2,11 @@ """Used to test Active Node functions.""" import pytest -from primaite.common.enums import FileSystemState, HardwareState, SoftwareState -from primaite.nodes.active_node import ActiveNode +# from primaite.common.enums import FileSystemState, HardwareState, SoftwareState +# from primaite.nodes.active_node import ActiveNode +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ @@ -36,6 +37,7 @@ def test_os_state_change(operating_state, expected_state): assert active_node.software_state == expected_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ @@ -66,6 +68,7 @@ def test_os_state_change_if_not_compromised(operating_state, expected_state): assert active_node.software_state == expected_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ @@ -92,6 +95,7 @@ def test_file_system_change(operating_state, expected_state): assert active_node.file_system_state_actual == expected_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ diff --git a/tests/test_full_legacy_config_session.py b/tests/test_full_legacy_config_session.py index 066ff72c..ac727a22 100644 --- a/tests/test_full_legacy_config_session.py +++ b/tests/test_full_legacy_config_session.py @@ -6,6 +6,7 @@ from primaite.main import run from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "legacy_file", [ diff --git a/tests/test_lay_down_config.py b/tests/test_lay_down_config.py index 99e66708..83c6063c 100644 --- a/tests/test_lay_down_config.py +++ b/tests/test_lay_down_config.py @@ -2,16 +2,17 @@ import pytest import yaml -from primaite.config.lay_down_config import ( - convert_legacy_lay_down_config, - data_manipulation_config_path, - ddos_basic_one_config_path, - ddos_basic_two_config_path, - dos_very_basic_config_path, -) +# from primaite.config.lay_down_config import ( +# convert_legacy_lay_down_config, +# data_manipulation_config_path, +# ddos_basic_one_config_path, +# ddos_basic_two_config_path, +# dos_very_basic_config_path, +# ) from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "legacy_file, new_path", [ diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index ff3528e1..b138dd5e 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -4,10 +4,11 @@ import numpy as np import pytest -from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler +# from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ @@ -29,6 +30,7 @@ def test_default_obs_space(temp_primaite_session): assert isinstance(components[0], NodeLinkTable) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ @@ -51,6 +53,7 @@ def test_registering_components(temp_primaite_session): assert component not in handler.registered_obs_components +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ @@ -140,6 +143,7 @@ class TestNodeLinkTable: ) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ @@ -194,6 +198,7 @@ class TestNodeStatuses: assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0]) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ @@ -240,6 +245,7 @@ class TestLinkTrafficLevels: assert np.array_equal(obs, [6, 0, 6, 0]) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ diff --git a/tests/test_primaite_session.py b/tests/test_primaite_session.py index 6e23b3ac..bf9332a7 100644 --- a/tests/test_primaite_session.py +++ b/tests/test_primaite_session.py @@ -4,12 +4,14 @@ import os import pytest from primaite import getLogger -from primaite.config.lay_down_config import dos_very_basic_config_path + +# from primaite.config.lay_down_config import dos_very_basic_config_path from tests import TEST_CONFIG_ROOT _LOGGER = getLogger(__name__) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ diff --git a/tests/test_red_random_agent_behaviour.py b/tests/test_red_random_agent_behaviour.py index e99f4adb..3f999f9b 100644 --- a/tests/test_red_random_agent_behaviour.py +++ b/tests/test_red_random_agent_behaviour.py @@ -1,11 +1,12 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK import pytest -from primaite.config.lay_down_config import data_manipulation_config_path -from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed +# from primaite.config.lay_down_config import data_manipulation_config_path +# from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ diff --git a/tests/test_resetting_node.py b/tests/test_resetting_node.py index d4e27c17..53779c4f 100644 --- a/tests/test_resetting_node.py +++ b/tests/test_resetting_node.py @@ -2,13 +2,14 @@ """Used to test Active Node functions.""" import pytest -from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState -from primaite.common.service import Service -from primaite.config.training_config import TrainingConfig -from primaite.nodes.active_node import ActiveNode -from primaite.nodes.service_node import ServiceNode +# from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState +# from primaite.common.service import Service +# from primaite.config.training_config import TrainingConfig +# from primaite.nodes.active_node import ActiveNode +# from primaite.nodes.service_node import ServiceNode +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "starting_operating_state, expected_operating_state", [(HardwareState.RESETTING, HardwareState.ON)], @@ -35,6 +36,7 @@ def test_node_resets_correctly(starting_operating_state, expected_operating_stat assert active_node.hardware_state == expected_operating_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_operating_state", [(HardwareState.BOOTING, HardwareState.ON)], @@ -62,6 +64,7 @@ def test_node_boots_correctly(operating_state, expected_operating_state): assert service_node.hardware_state == expected_operating_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_operating_state", [(HardwareState.SHUTTING_DOWN, HardwareState.OFF)], diff --git a/tests/test_reward.py b/tests/test_reward.py index 2ac66af1..7a980d32 100644 --- a/tests/test_reward.py +++ b/tests/test_reward.py @@ -7,6 +7,7 @@ from tests import TEST_CONFIG_ROOT _LOGGER = getLogger(__name__) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [ diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index aff5496a..eed6f4d6 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -1,10 +1,11 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK import pytest as pytest -from primaite.config.lay_down_config import dos_very_basic_config_path +# from primaite.config.lay_down_config import dos_very_basic_config_path from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], @@ -49,6 +50,7 @@ def test_seeded_learning(temp_primaite_session): assert actual_mean_reward_per_episode == expected_mean_reward_per_episode +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "temp_primaite_session", [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], diff --git a/tests/test_service_node.py b/tests/test_service_node.py index 906bcf55..6b1bd1ee 100644 --- a/tests/test_service_node.py +++ b/tests/test_service_node.py @@ -2,11 +2,12 @@ """Used to test Service Node functions.""" import pytest -from primaite.common.enums import HardwareState, SoftwareState -from primaite.common.service import Service -from primaite.nodes.service_node import ServiceNode +# from primaite.common.enums import HardwareState, SoftwareState +# from primaite.common.service import Service +# from primaite.nodes.service_node import ServiceNode +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ @@ -39,6 +40,7 @@ def test_service_state_change(operating_state, expected_state): assert service_node.get_service_state("TCP") == expected_state +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.mark.parametrize( "operating_state, expected_state", [ diff --git a/tests/test_session_loading.py b/tests/test_session_loading.py index f9990f76..fdec4ede 100644 --- a/tests/test_session_loading.py +++ b/tests/test_session_loading.py @@ -6,14 +6,18 @@ from pathlib import Path from typing import Union from uuid import uuid4 +import pytest from typer.testing import CliRunner from primaite import getLogger -from primaite.agents.sb3 import SB3Agent + +# from primaite.agents.sb3 import SB3Agent from primaite.cli import app -from primaite.common.enums import AgentFramework, AgentIdentifier + +# from primaite.common.enums import AgentFramework, AgentIdentifier from primaite.main import run -from primaite.primaite_session import PrimaiteSession + +# from primaite.primaite_session import PrimaiteSession from primaite.utils.session_output_reader import av_rewards_dict from tests import TEST_ASSETS_ROOT @@ -62,6 +66,7 @@ def copy_session_asset(asset_path: Union[str, Path]) -> str: return copy_path +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_load_sb3_session(): """Test that loading an SB3 agent works.""" test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") @@ -104,6 +109,7 @@ def test_load_sb3_session(): shutil.rmtree(test_path) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_load_primaite_session(): """Test that loading a Primaite session works.""" test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") @@ -150,6 +156,7 @@ def test_load_primaite_session(): shutil.rmtree(test_path) +@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 def test_run_loading(): """Test loading session via main.run.""" test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 5d300232..949898cf 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -3,12 +3,13 @@ import time import pytest -from primaite.acl.acl_rule import ACLRule -from primaite.common.enums import HardwareState -from primaite.environment.primaite_env import Primaite +# from primaite.acl.acl_rule import ACLRule +# from primaite.common.enums import HardwareState +# from primaite.environment.primaite_env import Primaite from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") def run_generic_set_actions(env: Primaite): """Run against a generic agent with specified blue agent actions.""" # Reset the environment at the start of the episode @@ -47,6 +48,7 @@ def run_generic_set_actions(env: Primaite): # env.close() +@pytest.skip("Deprecated") @pytest.mark.parametrize( "temp_primaite_session", [ @@ -88,6 +90,7 @@ def test_single_action_space_is_valid(temp_primaite_session): assert both_action_spaces +@pytest.skip("Deprecated") @pytest.mark.parametrize( "temp_primaite_session", [ diff --git a/tests/test_train_eval_episode_steps.py b/tests/test_train_eval_episode_steps.py index 1b53fe9d..0be968ae 100644 --- a/tests/test_train_eval_episode_steps.py +++ b/tests/test_train_eval_episode_steps.py @@ -2,12 +2,14 @@ import pytest from primaite import getLogger -from primaite.config.lay_down_config import dos_very_basic_config_path + +# from primaite.config.lay_down_config import dos_very_basic_config_path from tests import TEST_CONFIG_ROOT _LOGGER = getLogger(__name__) +@pytest.skip("Deprecated") @pytest.mark.parametrize( "temp_primaite_session", [[TEST_CONFIG_ROOT / "train_episode_step.yaml", dos_very_basic_config_path()]], diff --git a/tests/test_training_config.py b/tests/test_training_config.py index 58f9c797..e4e3fa32 100644 --- a/tests/test_training_config.py +++ b/tests/test_training_config.py @@ -1,10 +1,12 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK +import pytest import yaml -from primaite.config import training_config +# from primaite.config import training_config from tests import TEST_CONFIG_ROOT +@pytest.skip("Deprecated") def test_legacy_lay_down_config_yaml_conversion(): """Tests the conversion of legacy lay down config files.""" legacy_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml" @@ -22,6 +24,7 @@ def test_legacy_lay_down_config_yaml_conversion(): assert converted_dict[key] == value +@pytest.skip("Deprecated") def test_create_config_values_main_from_file(): """Tests creating an instance of TrainingConfig from file.""" new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "new_training_config.yaml" @@ -29,6 +32,7 @@ def test_create_config_values_main_from_file(): training_config.load(new_path) +@pytest.skip("Deprecated") def test_create_config_values_main_from_legacy_file(): """Tests creating an instance of TrainingConfig from legacy file.""" new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml" From 2b8462d38da298b63e391f17994d2fccfac2376f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:04:04 +0100 Subject: [PATCH 36/53] Fix link import in v3 code --- src/primaite/simulator/network/hardware/nodes/switch.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/primaite/simulator/network/hardware/nodes/switch.py b/src/primaite/simulator/network/hardware/nodes/switch.py index 09b53483..fe61509c 100644 --- a/src/primaite/simulator/network/hardware/nodes/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/switch.py @@ -4,8 +4,7 @@ from prettytable import MARKDOWN, PrettyTable from primaite import getLogger from primaite.exceptions import NetworkError -from primaite.links.link import Link -from primaite.simulator.network.hardware.base import Node, SwitchPort +from primaite.simulator.network.hardware.base import Link, Node, SwitchPort from primaite.simulator.network.transmission.data_link_layer import Frame _LOGGER = getLogger(__name__) From 1e811148ed597bfe77793af5892f54a00622375f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:04:30 +0100 Subject: [PATCH 37/53] Fix dependency versions to align with GATE --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d7e71a28..51ed84f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ - "gym==0.21.0", + "gymnasium==0.28.1", "jupyterlab==3.6.1", "kaleido==0.2.1", "matplotlib==3.7.1", @@ -35,7 +35,7 @@ dependencies = [ "polars==0.18.4", "prettytable==3.8.0", "PyYAML==6.0", - "stable-baselines3==1.6.2", + "stable-baselines3[extra]==2.1.0", "tensorflow==2.12.0", "typer[all]==0.9.0", "pydantic==2.1.1" From c57b5152c01f8d435755370d0ef8b1e02586efda Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:08:01 +0100 Subject: [PATCH 38/53] Remove broken tests --- tests/conftest.py | 40 +- .../game_layer/test_observations.py | 2 +- tests/test_acl.py | 174 -------- tests/test_active_node.py | 126 ------ tests/test_full_legacy_config_session.py | 30 -- tests/test_lay_down_config.py | 45 -- tests/test_observation_space.py | 383 ------------------ tests/test_primaite_session.py | 79 ---- tests/test_red_random_agent_behaviour.py | 40 -- tests/test_resetting_node.py | 89 ---- tests/test_reward.py | 54 --- .../test_seeding_and_deterministic_session.py | 66 --- tests/test_service_node.py | 73 ---- tests/test_session_loading.py | 194 --------- tests/test_single_action_space.py | 132 ------ tests/test_train_eval_episode_steps.py | 45 -- tests/test_training_config.py | 40 -- 17 files changed, 21 insertions(+), 1591 deletions(-) delete mode 100644 tests/test_acl.py delete mode 100644 tests/test_active_node.py delete mode 100644 tests/test_full_legacy_config_session.py delete mode 100644 tests/test_lay_down_config.py delete mode 100644 tests/test_observation_space.py delete mode 100644 tests/test_primaite_session.py delete mode 100644 tests/test_red_random_agent_behaviour.py delete mode 100644 tests/test_resetting_node.py delete mode 100644 tests/test_reward.py delete mode 100644 tests/test_seeding_and_deterministic_session.py delete mode 100644 tests/test_service_node.py delete mode 100644 tests/test_session_loading.py delete mode 100644 tests/test_single_action_space.py delete mode 100644 tests/test_train_eval_episode_steps.py delete mode 100644 tests/test_training_config.py diff --git a/tests/conftest.py b/tests/conftest.py index 425acc09..2701955c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -54,37 +54,37 @@ def file_system() -> FileSystem: return Node(hostname="fs_node").file_system -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 # PrimAITE v2 stuff -class TempPrimaiteSession(PrimaiteSession): +@pytest.mark.skip("Deprecated") # TODO: implement a similar test for primaite v3 +class TempPrimaiteSession: # PrimaiteSession): """ A temporary PrimaiteSession class. Uses context manager for deletion of files upon exit. """ - def __init__( - self, - training_config_path: Union[str, Path], - lay_down_config_path: Union[str, Path], - ): - super().__init__(training_config_path, lay_down_config_path) - self.setup() + # def __init__( + # self, + # training_config_path: Union[str, Path], + # lay_down_config_path: Union[str, Path], + # ): + # super().__init__(training_config_path, lay_down_config_path) + # self.setup() - @property - def env(self) -> Primaite: - """Direct access to the env for ease of testing.""" - return self._agent_session._env # noqa + # @property + # def env(self) -> Primaite: + # """Direct access to the env for ease of testing.""" + # return self._agent_session._env # noqa - def __enter__(self): - return self + # def __enter__(self): + # return self - def __exit__(self, type, value, tb): - shutil.rmtree(self.session_path) - _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") + # def __exit__(self, type, value, tb): + # shutil.rmtree(self.session_path) + # _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 +@pytest.mark.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.fixture def temp_primaite_session(request): """ @@ -139,7 +139,7 @@ def temp_primaite_session(request): return TempPrimaiteSession(training_config_path, lay_down_config_path) -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 +@pytest.mark.skip("Deprecated") # TODO: implement a similar test for primaite v3 @pytest.fixture def temp_session_path() -> Path: """ diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index c1f20d78..97154f62 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -1,4 +1,4 @@ -from gym import spaces +from gymnasium import spaces from primaite.game.agent.observations import FileObservation from primaite.simulator.network.hardware.nodes.computer import Computer diff --git a/tests/test_acl.py b/tests/test_acl.py deleted file mode 100644 index a0dfb997..00000000 --- a/tests/test_acl.py +++ /dev/null @@ -1,174 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Used to tes the ACL functions.""" - -# from primaite.acl.access_control_list import AccessControlList -# from primaite.acl.acl_rule import ACLRule -# from primaite.common.enums import RulePermissionType - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_acl_address_match_1(): - """Test that matching IP addresses produce True.""" - acl = AccessControlList(RulePermissionType.DENY, 10) - - rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "192.168.1.2", "TCP", "80") - - assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_acl_address_match_2(): - """Test that mismatching IP addresses produce False.""" - acl = AccessControlList(RulePermissionType.DENY, 10) - - rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "192.168.1.2", "TCP", "80") - - assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_acl_address_match_3(): - """Test the ANY condition for source IP addresses produce True.""" - acl = AccessControlList(RulePermissionType.DENY, 10) - - rule = ACLRule(RulePermissionType.ALLOW, "ANY", "192.168.1.2", "TCP", "80") - - assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_acl_address_match_4(): - """Test the ANY condition for dest IP addresses produce True.""" - acl = AccessControlList(RulePermissionType.DENY, 10) - - rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "ANY", "TCP", "80") - - assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_check_acl_block_affirmative(): - """Test the block function (affirmative).""" - # Create the Access Control List - acl = AccessControlList(RulePermissionType.DENY, 10) - - # Create a rule - acl_rule_permission = RulePermissionType.ALLOW - acl_rule_source = "192.168.1.1" - acl_rule_destination = "192.168.1.2" - acl_rule_protocol = "TCP" - acl_rule_port = "80" - acl_position_in_list = "0" - - acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_position_in_list, - ) - assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_check_acl_block_negative(): - """Test the block function (negative).""" - # Create the Access Control List - acl = AccessControlList(RulePermissionType.DENY, 10) - - # Create a rule - acl_rule_permission = RulePermissionType.DENY - acl_rule_source = "192.168.1.1" - acl_rule_destination = "192.168.1.2" - acl_rule_protocol = "TCP" - acl_rule_port = "80" - acl_position_in_list = "0" - - acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_position_in_list, - ) - - assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == True - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_rule_hash(): - """Test the rule hash.""" - # Create the Access Control List - acl = AccessControlList(RulePermissionType.DENY, 10) - - rule = ACLRule(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80") - hash_value_local = hash(rule) - - hash_value_remote = acl.get_dictionary_hash(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80") - - assert hash_value_local == hash_value_remote - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_delete_rule(): - """Adds 3 rules and deletes 1 rule and checks its deletion.""" - # Create the Access Control List - acl = AccessControlList(RulePermissionType.ALLOW, 10) - - # Create a first rule - acl_rule_permission = RulePermissionType.DENY - acl_rule_source = "192.168.1.1" - acl_rule_destination = "192.168.1.2" - acl_rule_protocol = "TCP" - acl_rule_port = "80" - acl_position_in_list = "0" - - acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_position_in_list, - ) - - # Create a second rule - acl_rule_permission = RulePermissionType.DENY - acl_rule_source = "20" - acl_rule_destination = "30" - acl_rule_protocol = "FTP" - acl_rule_port = "21" - acl_position_in_list = "2" - - acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_position_in_list, - ) - - # Create a third rule - acl_rule_permission = RulePermissionType.ALLOW - acl_rule_source = "192.168.1.3" - acl_rule_destination = "192.168.1.1" - acl_rule_protocol = "UDP" - acl_rule_port = "60" - acl_position_in_list = "4" - - acl.add_rule( - acl_rule_permission, - acl_rule_source, - acl_rule_destination, - acl_rule_protocol, - acl_rule_port, - acl_position_in_list, - ) - # Remove the second ACL rule added from the list - acl.remove_rule(RulePermissionType.DENY, "20", "30", "FTP", "21") - - assert len(acl.acl) == 10 - assert acl.acl[2] is None diff --git a/tests/test_active_node.py b/tests/test_active_node.py deleted file mode 100644 index cf532bb8..00000000 --- a/tests/test_active_node.py +++ /dev/null @@ -1,126 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Used to test Active Node functions.""" -import pytest - -# from primaite.common.enums import FileSystemState, HardwareState, SoftwareState -# from primaite.nodes.active_node import ActiveNode - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, SoftwareState.GOOD), - (HardwareState.ON, SoftwareState.OVERWHELMED), - ], -) -def test_os_state_change(operating_state, expected_state): - """ - Test that a node cannot change its Software State. - - When its hardware state is OFF. - """ - active_node = ActiveNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - SoftwareState.GOOD, - "GOOD", - 1, - ) - - active_node.software_state = SoftwareState.OVERWHELMED - - assert active_node.software_state == expected_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, SoftwareState.GOOD), - (HardwareState.ON, SoftwareState.OVERWHELMED), - ], -) -def test_os_state_change_if_not_compromised(operating_state, expected_state): - """ - Test that a node cannot change its Software State. - - If not compromised) when its hardware state is OFF. - """ - active_node = ActiveNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - SoftwareState.GOOD, - "GOOD", - 1, - ) - - active_node.set_software_state_if_not_compromised(SoftwareState.OVERWHELMED) - - assert active_node.software_state == expected_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, FileSystemState.GOOD), - (HardwareState.ON, FileSystemState.CORRUPT), - ], -) -def test_file_system_change(operating_state, expected_state): - """Test that a node cannot change its file system state when its hardware state is ON.""" - active_node = ActiveNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - "COMPROMISED", - FileSystemState.GOOD, - 1, - ) - - active_node.set_file_system_state(FileSystemState.CORRUPT) - - assert active_node.file_system_state_actual == expected_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, FileSystemState.GOOD), - (HardwareState.ON, FileSystemState.CORRUPT), - ], -) -def test_file_system_change_if_not_compromised(operating_state, expected_state): - """ - Test that a node cannot change its file system state. - - If not compromised) when its hardware state is OFF. - """ - active_node = ActiveNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - "GOOD", - FileSystemState.GOOD, - 1, - ) - - active_node.set_file_system_state_if_not_compromised(FileSystemState.CORRUPT) - - assert active_node.file_system_state_actual == expected_state diff --git a/tests/test_full_legacy_config_session.py b/tests/test_full_legacy_config_session.py deleted file mode 100644 index ac727a22..00000000 --- a/tests/test_full_legacy_config_session.py +++ /dev/null @@ -1,30 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -import pytest - -from primaite.main import run -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "legacy_file", - [ - ("legacy_config_1_DDOS_BASIC.yaml"), - ("legacy_config_2_DDOS_BASIC.yaml"), - ("legacy_config_3_DOS_VERY_BASIC.yaml"), - ("legacy_config_5_DATA_MANIPULATION.yaml"), - ], -) -def test_legacy_training_config_run_session(legacy_file): - """Tests using legacy training and lay down config files in PrimAITE session end-to-end.""" - legacy_training_config_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml" - legacy_lay_down_config_path = TEST_CONFIG_ROOT / "legacy_conversion" / legacy_file - - # Run a PrimAITE session using legacy training and lay down config file paths - run( - legacy_training_config_path, - legacy_lay_down_config_path, - legacy_training_config=True, - legacy_lay_down_config=True, - ) diff --git a/tests/test_lay_down_config.py b/tests/test_lay_down_config.py deleted file mode 100644 index 83c6063c..00000000 --- a/tests/test_lay_down_config.py +++ /dev/null @@ -1,45 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest -import yaml - -# from primaite.config.lay_down_config import ( -# convert_legacy_lay_down_config, -# data_manipulation_config_path, -# ddos_basic_one_config_path, -# ddos_basic_two_config_path, -# dos_very_basic_config_path, -# ) -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "legacy_file, new_path", - [ - ("legacy_config_1_DDOS_BASIC.yaml", ddos_basic_one_config_path()), - ("legacy_config_2_DDOS_BASIC.yaml", ddos_basic_two_config_path()), - ("legacy_config_3_DOS_VERY_BASIC.yaml", dos_very_basic_config_path()), - ("legacy_config_5_DATA_MANIPULATION.yaml", data_manipulation_config_path()), - ], -) -def test_legacy_lay_down_config_load(legacy_file, new_path): - """Tests converting legacy lay down files into the new format.""" - with open(TEST_CONFIG_ROOT / "legacy_conversion" / legacy_file, "r") as file: - legacy_lay_down_config = yaml.safe_load(file) - - with open(new_path, "r") as file: - new_lay_down_config = yaml.safe_load(file) - - converted_lay_down_config = convert_legacy_lay_down_config(legacy_lay_down_config) - - assert len(converted_lay_down_config) == len(new_lay_down_config) - - for i, new_item in enumerate(new_lay_down_config): - converted_item = converted_lay_down_config[i] - - for key, val in new_item.items(): - if key == "position": - continue - assert key in converted_item - - assert val == converted_item[key] diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py deleted file mode 100644 index b138dd5e..00000000 --- a/tests/test_observation_space.py +++ /dev/null @@ -1,383 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Test env creation and behaviour with different observation spaces.""" - -import numpy as np -import pytest - -# from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", - ] - ], - indirect=True, -) -def test_default_obs_space(temp_primaite_session): - """Create environment with no obs space defined in config and check that the default obs space was created.""" - with temp_primaite_session as session: - session.env.update_environent_obs() - - components = session.env.obs_handler.registered_obs_components - - assert len(components) == 1 - assert isinstance(components[0], NodeLinkTable) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", - ] - ], - indirect=True, -) -def test_registering_components(temp_primaite_session): - """Test regitering and deregistering a component.""" - with temp_primaite_session as session: - env = session.env - handler = ObservationsHandler() - component = NodeStatuses(env) - handler.register(component) - assert component in handler.registered_obs_components - handler.deregister(component) - assert component not in handler.registered_obs_components - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_LINK_TABLE.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", - ] - ], - indirect=True, -) -class TestNodeLinkTable: - """Test the NodeLinkTable observation component (in isolation).""" - - def test_obs_shape(self, temp_primaite_session): - """Try creating env with box observation space.""" - with temp_primaite_session as session: - env = session.env - env.update_environent_obs() - - # we have three nodes and two links, with two service - # therefore the box observation space will have: - # * 5 rows (3 nodes + 2 links) - # * 6 columns (four fixed and two for the services) - assert env.env_obs.shape == (5, 6) - - def test_value(self, temp_primaite_session): - """ - Test that the observation is generated correctly. - - The laydown has: - * 3 nodes (2 service nodes and 1 active node) - * 2 services - * 2 links - - Both nodes have both services, and all states are GOOD, therefore the expected observation value is: - - * Node 1: - * 1 (id) - * 1 (good hardware state) - * 3 (compromised OS state) - * 1 (good file system state) - * 1 (good TCP state) - * 1 (good UDP state) - * Node 2: - * 2 (id) - * 1 (good hardware state) - * 1 (good OS state) - * 1 (good file system state) - * 1 (good TCP state) - * 4 (overwhelmed UDP state) - * Node 3 (active node): - * 3 (id) - * 1 (good hardware state) - * 1 (good OS state) - * 1 (good file system state) - * 0 (doesn't have service1) - * 0 (doesn't have service2) - * Link 1: - * 4 (id) - * 0 (n/a hardware state) - * 0 (n/a OS state) - * 0 (n/a file system state) - * 999 (999 traffic for service1) - * 0 (no traffic for service2) - * Link 2: - * 5 (id) - * 0 (good hardware state) - * 0 (good OS state) - * 0 (good file system state) - * 999 (999 traffic service1) - * 0 (no traffic for service2) - """ - with temp_primaite_session as session: - env = session.env - # act = np.asarray([0,]) - obs, reward, done, info = env.step(0) # apply the 'do nothing' action - - assert np.array_equal( - obs, - [ - [1, 1, 3, 1, 1, 1], - [2, 1, 1, 1, 1, 4], - [3, 1, 1, 1, 0, 0], - [4, 0, 0, 0, 999, 0], - [5, 0, 0, 0, 999, 0], - ], - ) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", - ] - ], - indirect=True, -) -class TestNodeStatuses: - """Test the NodeStatuses observation component (in isolation).""" - - def test_obs_shape(self, temp_primaite_session): - """Try creating env with NodeStatuses as the only component.""" - with temp_primaite_session as session: - env = session.env - assert env.env_obs.shape == (15,) - - def test_values(self, temp_primaite_session): - """ - Test that the hardware and software states are encoded correctly. - - The laydown has: - * one node with a compromised operating system state - * one node with two services, and the second service is overwhelmed. - * all other states are good or null - Therefore, the expected state is: - * node 1: - * hardware = good (1) - * OS = compromised (3) - * file system = good (1) - * service 1 = good (1) - * service 2 = good (1) - * node 2: - * hardware = good (1) - * OS = good (1) - * file system = good (1) - * service 1 = good (1) - * service 2 = overwhelmed (4) - * node 3 (switch): - * hardware = good (1) - * OS = good (1) - * file system = good (1) - * service 1 = n/a (0) - * service 2 = n/a (0) - """ - with temp_primaite_session as session: - env = session.env - obs, _, _, _ = env.step(0) # apply the 'do nothing' action - print(obs) - assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0]) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", - ] - ], - indirect=True, -) -class TestLinkTrafficLevels: - """Test the LinkTrafficLevels observation component (in isolation).""" - - def test_obs_shape(self, temp_primaite_session): - """Try creating env with MultiDiscrete observation space.""" - with temp_primaite_session as session: - env = session.env - env.update_environent_obs() - - # we have two links and two services, so the shape should be 2 * 2 - assert env.env_obs.shape == (2 * 2,) - - def test_values(self, temp_primaite_session): - """ - Test that traffic values are encoded correctly. - - The laydown has: - * two services - * three nodes - * two links - * an IER trying to send 999 bits of data over both links the whole time (via the first service) - * link bandwidth of 1000, therefore the utilisation is 99.9% - """ - with temp_primaite_session as session: - env = session.env - obs, reward, done, info = env.step(0) - obs, reward, done, info = env.step(0) - - # the observation space has combine_service_traffic set to False, so the space has this format: - # [link1_service1, link1_service2, link2_service1, link2_service2] - # we send 999 bits of data via link1 and link2 on service 1. - # therefore the first and third elements should be 6 and all others 0 - # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) - assert np.array_equal(obs, [6, 0, 6, 0]) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "obs_tests/main_config_ACCESS_CONTROL_LIST.yaml", - TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml", - ] - ], - indirect=True, -) -class TestAccessControlList: - """Test the AccessControlList observation component (in isolation).""" - - def test_obs_shape(self, temp_primaite_session): - """Try creating env with MultiDiscrete observation space. - - The laydown has 3 ACL Rules - that is the maximum_acl_rules it can have. - Each ACL Rule in the observation space has 6 different elements: - - 6 * 3 = 18 - """ - with temp_primaite_session as session: - env = session.env - env.update_environent_obs() - - assert env.env_obs.shape == (18,) - - def test_values(self, temp_primaite_session): - """Test that traffic values are encoded correctly. - - The laydown has: - * one ACL IMPLICIT DENY rule - - Therefore, the ACL is full of NAs aka zeros and just 6 non-zero elements representing DENY ANY ANY ANY at - Position 2. - """ - with temp_primaite_session as session: - env = session.env - obs, reward, done, info = env.step(0) - obs, reward, done, info = env.step(0) - - assert np.array_equal(obs, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2]) - - def test_observation_space_with_implicit_rule(self, temp_primaite_session): - """ - Test observation space is what is expected when an agent adds ACLs during an episode. - - At the start of the episode, there is a single implicit DENY rule - In the observation space IMPLICIT DENY: 1,1,1,1,1,0 - 0 shows the rule is the start (when episode began no other rules were created) so this is correct. - - On Step 2, there is an ACL rule added at Position 0: 2,2,3,2,3,0 - - On Step 4, there is a second ACL rule added at POSITION 1: 2,4,2,3,3,1 - - The final observation space should be this: - [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] - - The ACL Rule from Step 2 is added first and has a HIGHER position than the ACL rule from Step 4 - but both come before the IMPLICIT DENY which will ALWAYS be at the end of the ACL List. - """ - # TODO: Refactor this at some point to build a custom ACL Hardcoded - # Agent and then patch the AgentIdentifier Enum class so that it - # has ACL_AGENT. This then allows us to set the agent identified in - # the main config and is a bit cleaner. - - with temp_primaite_session as session: - env = session.env - training_config = env.training_config - for episode in range(0, training_config.num_train_episodes): - for step in range(0, training_config.num_train_steps): - # Do nothing action - action = 0 - if step == 2: - # Action to add the first ACL rule - action = 43 - elif step == 4: - # Action to add the second ACL rule - action = 96 - - # Run the simulation step on the live environment - obs, reward, done, info = env.step(action) - - # Break if done is True - if done: - break - obs = env.env_obs - - assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) - - def test_observation_space_with_different_positions(self, temp_primaite_session): - """ - Test observation space is what is expected when an agent adds ACLs during an episode. - - At the start of the episode, there is a single implicit DENY rule - In the observation space IMPLICIT DENY: 1,1,1,1,1,0 - 0 shows the rule is the start (when episode began no other rules were created) so this is correct. - - On Step 2, there is an ACL rule added at Position 1: 2,2,3,2,3,1 - - On Step 4 there is a second ACL rule added at Position 0: 2,4,2,3,3,0 - - The final observation space should be this: - [2 , 4, 2, 3, 3, 0, 2, 2, 3, 2, 3, 1, 1, 1, 1, 1, 1, 2] - - The ACL Rule from Step 2 is added before and has a LOWER position than the ACL rule from Step 4 - but both come before the IMPLICIT DENY which will ALWAYS be at the end of the ACL List. - """ - # TODO: Refactor this at some point to build a custom ACL Hardcoded - # Agent and then patch the AgentIdentifier Enum class so that it - # has ACL_AGENT. This then allows us to set the agent identified in - # the main config and is a bit cleaner. - - with temp_primaite_session as session: - env = session.env - training_config = env.training_config - for episode in range(0, training_config.num_train_episodes): - for step in range(0, training_config.num_train_steps): - # Do nothing action - action = 0 - if step == 2: - # Action to add the first ACL rule - action = 44 - elif step == 4: - # Action to add the second ACL rule - action = 95 - # Run the simulation step on the live environment - obs, reward, done, info = env.step(action) - - # Break if done is True - if done: - break - obs = env.env_obs - - assert np.array_equal(obs, [2, 4, 2, 3, 3, 0, 2, 2, 3, 2, 3, 1, 1, 1, 1, 1, 1, 2]) diff --git a/tests/test_primaite_session.py b/tests/test_primaite_session.py deleted file mode 100644 index bf9332a7..00000000 --- a/tests/test_primaite_session.py +++ /dev/null @@ -1,79 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import os - -import pytest - -from primaite import getLogger - -# from primaite.config.lay_down_config import dos_very_basic_config_path -from tests import TEST_CONFIG_ROOT - -_LOGGER = getLogger(__name__) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - # [TEST_CONFIG_ROOT / "session_test/training_config_main_rllib.yaml", dos_very_basic_config_path()], - [TEST_CONFIG_ROOT / "session_test/training_config_main_sb3.yaml", dos_very_basic_config_path()], - ], - indirect=True, -) -def test_primaite_session(temp_primaite_session): - """ - Tests the PrimaiteSession class and all of its outputs. - - This test runs for both a Stable Baselines3 agent, and a Ray RLlib agent. - """ - with temp_primaite_session as session: - session_path = session.session_path - assert session_path.exists() - session.learn() - # Learning outputs are saved in session.learning_path - session.evaluate() - # Evaluation outputs are saved in session.evaluation_path - - # If you need to inspect any session outputs, it must be done inside - # the context manager - - # Check that the metadata json file exists - assert (session_path / "session_metadata.json").exists() - - # Check that the network png file exists - assert (session_path / f"network_{session.timestamp_str}.png").exists() - - # Check that the saved agent exists - assert session._agent_session._saved_agent_path.exists() - - # Check that both the transactions and av reward csv files exist - for file in session.learning_path.iterdir(): - if file.suffix == ".csv": - assert "all_transactions" in file.name or "average_reward_per_episode" in file.name - - # Check that both the transactions and av reward csv files exist - for file in session.evaluation_path.iterdir(): - if file.suffix == ".csv": - assert "all_transactions" in file.name or "average_reward_per_episode" in file.name - - # Check that the average reward per episode plots exist - assert (session.learning_path / f"average_reward_per_episode_{session.timestamp_str}.png").exists() - assert (session.evaluation_path / f"average_reward_per_episode_{session.timestamp_str}.png").exists() - - # Check that the metadata has captured the correct number of learning and eval episodes and steps - assert len(session.learn_av_reward_per_episode_dict().keys()) == 10 - assert len(session.learn_all_transactions_dict().keys()) == 10 * 256 - - assert len(session.eval_av_reward_per_episode_dict().keys()) == 3 - assert len(session.eval_all_transactions_dict().keys()) == 3 * 256 - - _LOGGER.debug("Inspecting files in temp session path...") - for dir_path, dir_names, file_names in os.walk(session_path): - for file in file_names: - path = os.path.join(dir_path, file) - file_str = path.split(str(session_path))[-1] - _LOGGER.debug(f" {file_str}") - - # Now that we've exited the context manager, the session.session_path - # directory and its contents are deleted - assert not session_path.exists() diff --git a/tests/test_red_random_agent_behaviour.py b/tests/test_red_random_agent_behaviour.py deleted file mode 100644 index 3f999f9b..00000000 --- a/tests/test_red_random_agent_behaviour.py +++ /dev/null @@ -1,40 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest - -# from primaite.config.lay_down_config import data_manipulation_config_path -# from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "test_random_red_main_config.yaml", - data_manipulation_config_path(), - ] - ], - indirect=True, -) -def test_random_red_agent_behaviour(temp_primaite_session): - """Test that red agent POL is randomised each episode.""" - list_of_node_instructions = [] - - with temp_primaite_session as session: - session.evaluate() - list_of_node_instructions.append(session.env.red_node_pol) - - session.evaluate() - list_of_node_instructions.append(session.env.red_node_pol) - - # compare instructions to make sure that red instructions are truly random - for index, instruction in enumerate(list_of_node_instructions): - for key in list_of_node_instructions[index].keys(): - instruction: NodeStateInstructionRed = list_of_node_instructions[index][key] - print(f"run {index}") - print(f"{key} start step: {instruction.get_start_step()}") - print(f"{key} end step: {instruction.get_end_step()}") - print(f"{key} target node id: {instruction.get_target_node_id()}") - print("") - assert list_of_node_instructions[0].__ne__(list_of_node_instructions[1]) diff --git a/tests/test_resetting_node.py b/tests/test_resetting_node.py deleted file mode 100644 index 53779c4f..00000000 --- a/tests/test_resetting_node.py +++ /dev/null @@ -1,89 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Used to test Active Node functions.""" -import pytest - -# from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState -# from primaite.common.service import Service -# from primaite.config.training_config import TrainingConfig -# from primaite.nodes.active_node import ActiveNode -# from primaite.nodes.service_node import ServiceNode - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "starting_operating_state, expected_operating_state", - [(HardwareState.RESETTING, HardwareState.ON)], -) -def test_node_resets_correctly(starting_operating_state, expected_operating_state): - """Tests that a node resets correctly.""" - active_node = ActiveNode( - node_id="0", - name="node", - node_type=NodeType.COMPUTER, - priority=Priority.P1, - hardware_state=starting_operating_state, - ip_address="192.168.0.1", - software_state=SoftwareState.COMPROMISED, - file_system_state=FileSystemState.CORRUPT, - config_values=TrainingConfig(), - ) - - for x in range(5): - active_node.update_resetting_status() - - assert active_node.software_state == SoftwareState.GOOD - assert active_node.file_system_state_actual == FileSystemState.GOOD - assert active_node.hardware_state == expected_operating_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_operating_state", - [(HardwareState.BOOTING, HardwareState.ON)], -) -def test_node_boots_correctly(operating_state, expected_operating_state): - """Tests that a node boots correctly.""" - service_node = ServiceNode( - node_id=0, - name="node", - node_type="COMPUTER", - priority="1", - hardware_state=operating_state, - ip_address="192.168.0.1", - software_state=SoftwareState.GOOD, - file_system_state="GOOD", - config_values=1, - ) - service_attributes = Service(name="node", port="80", software_state=SoftwareState.COMPROMISED) - service_node.add_service(service_attributes) - - for x in range(5): - service_node.update_booting_status() - - assert service_attributes.software_state == SoftwareState.GOOD - assert service_node.hardware_state == expected_operating_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_operating_state", - [(HardwareState.SHUTTING_DOWN, HardwareState.OFF)], -) -def test_node_shutdown_correctly(operating_state, expected_operating_state): - """Tests that a node shutdown correctly.""" - active_node = ActiveNode( - node_id=0, - name="node", - node_type="COMPUTER", - priority="1", - hardware_state=operating_state, - ip_address="192.168.0.1", - software_state=SoftwareState.GOOD, - file_system_state="GOOD", - config_values=1, - ) - - for x in range(5): - active_node.update_shutdown_status() - - assert active_node.hardware_state == expected_operating_state diff --git a/tests/test_reward.py b/tests/test_reward.py deleted file mode 100644 index 7a980d32..00000000 --- a/tests/test_reward.py +++ /dev/null @@ -1,54 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest - -from primaite import getLogger -from tests import TEST_CONFIG_ROOT - -_LOGGER = getLogger(__name__) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml", - TEST_CONFIG_ROOT / "one_node_states_on_off_lay_down_config.yaml", - ] - ], - indirect=True, -) -def test_rewards_are_being_penalised_at_each_step_function( - temp_primaite_session, -): - """ - Test that hardware state is penalised at each step. - - When the initial state is OFF compared to reference state which is ON. - - The config 'one_node_states_on_off_lay_down_config.yaml' has 15 steps: - On different steps, the laydown config has Pattern of Life (PoLs) which change a state of the node's attribute. - For example, turning the nodes' file system state to CORRUPT from its original state GOOD. - As a result these are the following rewards are activated: - File System State: corrupt_should_be_good = -10 * 2 (on Steps 1 & 2) - Hardware State: off_should_be_on = -10 * 2 (on Steps 4 & 5) - Service State: compromised_should_be_good = -20 * 2 (on Steps 7 & 8) - Software State: compromised_should_be_good = -20 * 2 (on Steps 10 & 11) - - The Pattern of Life (PoLs) last for 2 steps, so the agent is penalised twice. - - Note: This test run inherits from conftest.py where the PrimAITE environment is ran and the blue agent is hard-coded - to do NOTHING on every step. - We use Pattern of Lifes (PoLs) to change the nodes states and display that the agent is being penalised on all steps - where the live network node differs from the network reference node. - - Total Reward: -10 + -10 + -10 + -10 + -20 + -20 + -20 + -20 = -120 - Step Count: 15 - - For the 4 steps where this occurs the average reward is: - Average Reward: -8 (-120 / 15) - """ - with temp_primaite_session as session: - session.evaluate() - ev_rewards = session.eval_av_reward_per_episode_dict() - assert ev_rewards[1] == -8.0 diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py deleted file mode 100644 index eed6f4d6..00000000 --- a/tests/test_seeding_and_deterministic_session.py +++ /dev/null @@ -1,66 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest as pytest - -# from primaite.config.lay_down_config import dos_very_basic_config_path -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], - indirect=True, -) -def test_seeded_learning(temp_primaite_session): - """ - Test running seeded learning produces the same output when ran twice. - - .. note:: - - If this is failing, the hard-coded expected_mean_reward_per_episode - from a pre-trained agent will probably need to be updated. If the - env changes and those changed how this agent is trained, chances are - the mean rewards are going to be different. - - Run the test, but print out the session.learn_av_reward_per_episode() - before comparing it. Then copy the printed dict and replace the - expected_mean_reward_per_episode with those values. The test should - now work. If not, then you've got a bug :). - """ - expected_mean_reward_per_episode = { - 1: -20.7421875, - 2: -19.82421875, - 3: -17.01171875, - 4: -19.08203125, - 5: -21.93359375, - 6: -20.21484375, - 7: -15.546875, - 8: -12.08984375, - 9: -17.59765625, - 10: -14.6875, - } - - with temp_primaite_session as session: - assert ( - session._training_config.seed == 67890 - ), "Expected output is based upon a agent that was trained with seed 67890" - session.learn() - actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict() - - assert actual_mean_reward_per_episode == expected_mean_reward_per_episode - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "temp_primaite_session", - [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], - indirect=True, -) -def test_deterministic_evaluation(temp_primaite_session): - """Test running deterministic evaluation gives same av eward per episode.""" - with temp_primaite_session as session: - # do stuff - session.learn() - session.evaluate() - eval_mean_reward = session.eval_av_reward_per_episode_dict() - assert len(set(eval_mean_reward.values())) == 1 diff --git a/tests/test_service_node.py b/tests/test_service_node.py deleted file mode 100644 index 6b1bd1ee..00000000 --- a/tests/test_service_node.py +++ /dev/null @@ -1,73 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Used to test Service Node functions.""" -import pytest - -# from primaite.common.enums import HardwareState, SoftwareState -# from primaite.common.service import Service -# from primaite.nodes.service_node import ServiceNode - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, SoftwareState.GOOD), - (HardwareState.ON, SoftwareState.OVERWHELMED), - ], -) -def test_service_state_change(operating_state, expected_state): - """ - Test that a node cannot change the state of a running service. - - When its hardware state is OFF. - """ - service_node = ServiceNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - "COMPROMISED", - "RESTORING", - 1, - ) - service = Service("TCP", 80, SoftwareState.GOOD) - service_node.add_service(service) - - service_node.set_service_state("TCP", SoftwareState.OVERWHELMED) - - assert service_node.get_service_state("TCP") == expected_state - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -@pytest.mark.parametrize( - "operating_state, expected_state", - [ - (HardwareState.OFF, SoftwareState.GOOD), - (HardwareState.ON, SoftwareState.OVERWHELMED), - ], -) -def test_service_state_change_if_not_comprised(operating_state, expected_state): - """ - Test that a node cannot change the state of a running service. - - If not compromised when its hardware state is ON. - """ - service_node = ServiceNode( - 0, - "node", - "COMPUTER", - "1", - operating_state, - "192.168.0.1", - "GOOD", - "RESTORING", - 1, - ) - service = Service("TCP", 80, SoftwareState.GOOD) - service_node.add_service(service) - - service_node.set_service_state_if_not_compromised("TCP", SoftwareState.OVERWHELMED) - - assert service_node.get_service_state("TCP") == expected_state diff --git a/tests/test_session_loading.py b/tests/test_session_loading.py deleted file mode 100644 index fdec4ede..00000000 --- a/tests/test_session_loading.py +++ /dev/null @@ -1,194 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import os.path -import shutil -import tempfile -from pathlib import Path -from typing import Union -from uuid import uuid4 - -import pytest -from typer.testing import CliRunner - -from primaite import getLogger - -# from primaite.agents.sb3 import SB3Agent -from primaite.cli import app - -# from primaite.common.enums import AgentFramework, AgentIdentifier -from primaite.main import run - -# from primaite.primaite_session import PrimaiteSession -from primaite.utils.session_output_reader import av_rewards_dict -from tests import TEST_ASSETS_ROOT - -_LOGGER = getLogger(__name__) - -runner = CliRunner() - -sb3_expected_avg_reward_per_episode = { - 10: 0.0, - 11: -0.0011074218750000008, - 12: -0.0010000000000000007, - 13: -0.0016601562500000013, - 14: -0.001400390625000001, - 15: -0.0009863281250000007, - 16: -0.0011855468750000008, - 17: -0.0009511718750000007, - 18: -0.0008789062500000007, - 19: -0.0012226562500000009, - 20: -0.0010292968750000007, -} - -sb3_expected_eval_rewards = -0.0018515625000000014 - - -def copy_session_asset(asset_path: Union[str, Path]) -> str: - """Copies the asset into a temporary test folder.""" - if asset_path is None: - raise Exception("No path provided") - - if isinstance(asset_path, Path): - asset_path = str(os.path.normpath(asset_path)) - - copy_path = str(Path(tempfile.gettempdir()) / "primaite" / str(uuid4())) - - # copy the asset into a temp path - try: - shutil.copytree(asset_path, copy_path) - except Exception as e: - msg = f"Unable to copy directory: {asset_path}" - _LOGGER.error(msg, e) - print(msg, e) - - _LOGGER.debug(f"Copied test asset to: {copy_path}") - - # return the copied assets path - return copy_path - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_load_sb3_session(): - """Test that loading an SB3 agent works.""" - test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") - - loaded_agent = SB3Agent(session_path=test_path) - - # loaded agent should have the same UUID as the previous agent - assert loaded_agent.uuid == "301874d3-2e14-43c2-ba7f-e2b03ad05dde" - assert loaded_agent._training_config.agent_framework == AgentFramework.SB3.name - assert loaded_agent._training_config.agent_identifier == AgentIdentifier.PPO.name - assert loaded_agent._training_config.deterministic - assert loaded_agent._training_config.seed == 12345 - assert str(loaded_agent.session_path) == str(test_path) - - # run another learn session - loaded_agent.learn() - - learn_mean_rewards = av_rewards_dict( - loaded_agent.learning_path / f"average_reward_per_episode_{loaded_agent.timestamp_str}.csv" - ) - - # run is seeded so should have the expected learn value - assert learn_mean_rewards == sb3_expected_avg_reward_per_episode - - # run an evaluation - loaded_agent.evaluate() - - # load the evaluation average reward csv file - eval_mean_reward = av_rewards_dict( - loaded_agent.evaluation_path / f"average_reward_per_episode_{loaded_agent.timestamp_str}.csv" - ) - - # the agent config ran the evaluation in deterministic mode, so should have the same reward value - assert len(set(eval_mean_reward.values())) == 1 - - # the evaluation should be the same as a previous run - assert next(iter(set(eval_mean_reward.values()))) == sb3_expected_eval_rewards - - # delete the test directory - shutil.rmtree(test_path) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_load_primaite_session(): - """Test that loading a Primaite session works.""" - test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") - - # create loaded session - session = PrimaiteSession(session_path=test_path) - - # run setup on session - session.setup() - - # make sure that the session was loaded correctly - assert session._agent_session.uuid == "301874d3-2e14-43c2-ba7f-e2b03ad05dde" - assert session._agent_session._training_config.agent_framework == AgentFramework.SB3.name - assert session._agent_session._training_config.agent_identifier == AgentIdentifier.PPO.name - assert session._agent_session._training_config.deterministic - assert session._agent_session._training_config.seed == 12345 - assert str(session._agent_session.session_path) == str(test_path) - - # run another learn session - session.learn() - - learn_mean_rewards = av_rewards_dict( - session.learning_path / f"average_reward_per_episode_{session.timestamp_str}.csv" - ) - - # run is seeded so should have the expected learn value - assert learn_mean_rewards == sb3_expected_avg_reward_per_episode - - # run an evaluation - session.evaluate() - - # load the evaluation average reward csv file - eval_mean_reward = av_rewards_dict( - session.evaluation_path / f"average_reward_per_episode_{session.timestamp_str}.csv" - ) - - # the agent config ran the evaluation in deterministic mode, so should have the same reward value - assert len(set(eval_mean_reward.values())) == 1 - - # the evaluation should be the same as a previous run - assert next(iter(set(eval_mean_reward.values()))) == sb3_expected_eval_rewards - - # delete the test directory - shutil.rmtree(test_path) - - -@pytest.skip("Deprecated") # TODO: implement a similar test for primaite v3 -def test_run_loading(): - """Test loading session via main.run.""" - test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") - - # create loaded session - run(session_path=test_path) - - learn_mean_rewards = av_rewards_dict( - next(Path(test_path).rglob("**/learning/average_reward_per_episode_*.csv"), None) - ) - - # run is seeded so should have the expected learn value - assert learn_mean_rewards == sb3_expected_avg_reward_per_episode - - # delete the test directory - shutil.rmtree(test_path) - - -def test_cli(): - """Test loading session via CLI.""" - test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session") - result = runner.invoke(app, ["session", "--load", test_path]) - - # cli should work - assert result.exit_code == 0 - - learn_mean_rewards = av_rewards_dict( - next(Path(test_path).rglob("**/learning/average_reward_per_episode_*.csv"), None) - ) - - # run is seeded so should have the expected learn value - assert learn_mean_rewards == sb3_expected_avg_reward_per_episode - - # delete the test directory - shutil.rmtree(test_path) diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py deleted file mode 100644 index 949898cf..00000000 --- a/tests/test_single_action_space.py +++ /dev/null @@ -1,132 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import time - -import pytest - -# from primaite.acl.acl_rule import ACLRule -# from primaite.common.enums import HardwareState -# from primaite.environment.primaite_env import Primaite -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") -def run_generic_set_actions(env: Primaite): - """Run against a generic agent with specified blue agent actions.""" - # Reset the environment at the start of the episode - # env.reset() - training_config = env.training_config - for episode in range(0, training_config.num_train_episodes): - for step in range(0, training_config.num_train_steps): - # Send the observation space to the agent to get an action - # TEMP - random action for now - # action = env.blue_agent_action(obs) - action = 0 - # print("Episode:", episode, "\nStep:", step) - if step == 5: - # [1, 1, 2, 1, 1, 1, 1(position)] - # Creates an ACL rule - # Allows traffic from server_1 to node_1 on port FTP - action = 56 - elif step == 7: - # [1, 1, 2, 0] Node Action - # Sets Node 1 Hardware State to OFF - # Does not resolve any service - action = 128 - # Run the simulation step on the live environment - obs, reward, done, info = env.step(action) - - # Break if done is True - if done: - break - - # Introduce a delay between steps - time.sleep(training_config.time_delay / 1000) - - # Reset the environment at the end of the episode - # env.reset() - - # env.close() - - -@pytest.skip("Deprecated") -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "single_action_space_main_config.yaml", - TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml", - ] - ], - indirect=True, -) -def test_single_action_space_is_valid(temp_primaite_session): - """Test single action space is valid.""" - # TODO: Refactor this at some point to build a custom ACL Hardcoded - # Agent and then patch the AgentIdentifier Enum class so that it - # has ACL_AGENT. This then allows us to set the agent identified in - # the main config and is a bit cleaner. - with temp_primaite_session as session: - env = session.env - - run_generic_set_actions(env) - # Retrieve the action space dictionary values from environment - env_action_space_dict = env.action_dict.values() - # Flags to check the conditions of the action space - contains_acl_actions = False - contains_node_actions = False - both_action_spaces = False - # Loop through each element of the list (which is every value from the dictionary) - for dict_item in env_action_space_dict: - # Node action detected - if len(dict_item) == 4: - contains_node_actions = True - # Link action detected - elif len(dict_item) == 7: - contains_acl_actions = True - # If both are there then the ANY action type is working - if contains_node_actions and contains_acl_actions: - both_action_spaces = True - # Check condition should be True - assert both_action_spaces - - -@pytest.skip("Deprecated") -@pytest.mark.parametrize( - "temp_primaite_session", - [ - [ - TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml", - TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml", - ] - ], - indirect=True, -) -def test_agent_is_executing_actions_from_both_spaces(temp_primaite_session): - """Test to ensure the blue agent is carrying out both kinds of operations (NODE & ACL).""" - # TODO: Refactor this at some point to build a custom ACL Hardcoded - # Agent and then patch the AgentIdentifier Enum class so that it - # has ACL_AGENT. This then allows us to set the agent identified in - # the main config and is a bit cleaner. - with temp_primaite_session as session: - env = session.env - # Run environment with specified fixed blue agent actions only - run_generic_set_actions(env) - # Retrieve hardware state of computer_1 node in laydown config - # Agent turned this off in Step 5 - computer_node_hardware_state = env.nodes["1"].hardware_state - # Retrieve the Access Control List object stored by the environment at the end of the episode - access_control_list = env.acl - # Use the Access Control List object acl object attribute to get dictionary - # Use dictionary.values() to get total list of all items in the dictionary - acl_rules_list = access_control_list.acl - # Length of this list tells you how many items are in the dictionary - # This number is the frequency of Access Control Rules in the environment - # In the scenario, we specified that the agent should create only 1 acl rule - # This 1 rule added to the implicit deny means there should be 2 rules in total. - rules_count = 0 - for rule in acl_rules_list: - if isinstance(rule, ACLRule): - rules_count += 1 - # Therefore these statements below MUST be true - assert computer_node_hardware_state == HardwareState.OFF - assert rules_count == 2 diff --git a/tests/test_train_eval_episode_steps.py b/tests/test_train_eval_episode_steps.py deleted file mode 100644 index 0be968ae..00000000 --- a/tests/test_train_eval_episode_steps.py +++ /dev/null @@ -1,45 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest - -from primaite import getLogger - -# from primaite.config.lay_down_config import dos_very_basic_config_path -from tests import TEST_CONFIG_ROOT - -_LOGGER = getLogger(__name__) - - -@pytest.skip("Deprecated") -@pytest.mark.parametrize( - "temp_primaite_session", - [[TEST_CONFIG_ROOT / "train_episode_step.yaml", dos_very_basic_config_path()]], - indirect=True, -) -def test_eval_steps_differ_from_training(temp_primaite_session): - """Uses PrimaiteSession class to compare number of episodes used for training and evaluation. - - Train_episode_step.yaml main config: - num_train_steps = 25 - num_train_episodes = 3 - num_eval_steps = 17 - num_eval_episodes = 1 - """ - expected_learning_metadata = {"total_episodes": 3, "total_time_steps": 75} - expected_evaluation_metadata = {"total_episodes": 1, "total_time_steps": 17} - - with temp_primaite_session as session: - # Run learning and check episode and step counts - session.learn() - assert session.env.actual_episode_count == expected_learning_metadata["total_episodes"] - assert session.env.total_step_count == expected_learning_metadata["total_time_steps"] - - # Run evaluation and check episode and step counts - session.evaluate() - assert session.env.actual_episode_count == expected_evaluation_metadata["total_episodes"] - assert session.env.total_step_count == expected_evaluation_metadata["total_time_steps"] - - # Load the session_metadata.json file and check that the both the - # learning and evaluation match what is expected above - metadata = session.metadata_file_as_dict() - assert metadata["learning"] == expected_learning_metadata - assert metadata["evaluation"] == expected_evaluation_metadata diff --git a/tests/test_training_config.py b/tests/test_training_config.py deleted file mode 100644 index e4e3fa32..00000000 --- a/tests/test_training_config.py +++ /dev/null @@ -1,40 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -import pytest -import yaml - -# from primaite.config import training_config -from tests import TEST_CONFIG_ROOT - - -@pytest.skip("Deprecated") -def test_legacy_lay_down_config_yaml_conversion(): - """Tests the conversion of legacy lay down config files.""" - legacy_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml" - new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "new_training_config.yaml" - - with open(legacy_path, "r") as file: - legacy_dict = yaml.safe_load(file) - - with open(new_path, "r") as file: - new_dict = yaml.safe_load(file) - - converted_dict = training_config.convert_legacy_training_config_dict(legacy_dict) - - for key, value in new_dict.items(): - assert converted_dict[key] == value - - -@pytest.skip("Deprecated") -def test_create_config_values_main_from_file(): - """Tests creating an instance of TrainingConfig from file.""" - new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "new_training_config.yaml" - - training_config.load(new_path) - - -@pytest.skip("Deprecated") -def test_create_config_values_main_from_legacy_file(): - """Tests creating an instance of TrainingConfig from legacy file.""" - new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml" - - training_config.load(new_path, legacy_file=True) From 4872c939ff46e53eaf31e6e67ba045b807343c3d Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:19:24 +0100 Subject: [PATCH 39/53] 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 78a8d2be3e531cd4f6d4872484a64048c9c16882 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 17:57:19 +0100 Subject: [PATCH 40/53] Fix File observation test --- src/primaite/simulator/file_system/file_system.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 30cda446..e4413313 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -92,8 +92,8 @@ class FileSystemItemABC(SimComponent): """ state = super().describe_state() state["name"] = self.name - state["status"] = self.health_status.value - state["visible_status"] = self.visible_health_status.value + state["health_status"] = self.health_status.value + state["visible_health_status"] = self.visible_health_status.value state["previous_hash"] = self.previous_hash return state From 38b71c0c8eef4aa5aefe122a137d6e5eeac5efcb Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 18:06:33 +0100 Subject: [PATCH 41/53] Update CHANGELOG --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c29c325a..3af5c14c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,12 @@ SessionManager. - FTP Services: `FTPClient` and `FTPServer` - HTTP Services: `WebBrowser` to simulate a web client and `WebServer` +### Removed +- Removed legacy simulation modules: `acl`, `common`, `environment`, `links`, `nodes`, `pol` +- Removed legacy training modules, they are replaced by the new ARCD GATE dependency +- Removed tests for legacy code + + ## [2.0.0] - 2023-07-26 ### Added From 02901a7c99e2bc5bfe7e9bf281c51a8aee746d84 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 19:07:45 +0100 Subject: [PATCH 42/53] 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 From afa7916db08f03407d40cf372bcd729fec69148a Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 25 Oct 2023 23:32:52 +0100 Subject: [PATCH 43/53] Updates to documentation --- docs/index.rst | 28 +- docs/source/about.rst | 704 ++++++++++++--------------- docs/source/action_system.rst | 88 ---- docs/source/config(v2).rst | 491 ------------------- docs/source/custom_agent.rst | 130 +---- docs/source/getting_started.rst | 60 ++- docs/source/migration_1.2_-_2.0.rst | 57 --- docs/source/primaite_session.rst | 359 +++++++------- docs/source/request_system.rst | 90 ++++ docs/source/simulation.rst | 6 +- docs/source/simulation_structure.rst | 6 +- docs/source/state_system.rst | 31 ++ 12 files changed, 673 insertions(+), 1377 deletions(-) delete mode 100644 docs/source/action_system.rst delete mode 100644 docs/source/config(v2).rst delete mode 100644 docs/source/migration_1.2_-_2.0.rst create mode 100644 docs/source/request_system.rst create mode 100644 docs/source/state_system.rst diff --git a/docs/index.rst b/docs/index.rst index ed66797d..fa877064 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -92,26 +92,34 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! .. toctree:: :maxdepth: 8 - :caption: Contents: + :caption: About PrimAITE: + :hidden: + + source/about + source/dependencies + source/glossary + +.. toctree:: + :caption: Usage: :hidden: source/getting_started - source/about - source/config - source/config(v2) + source/primaite_session source/simulation source/game_layer - source/primaite_session source/custom_agent + source/config + +.. toctree:: + :caption: Developer information: + :hidden: + + source/state_system + source/request_system PrimAITE API PrimAITE Tests - source/dependencies - source/glossary - source/migration_1.2_-_2.0 -.. TODO: Add project links once public repo has been created - .. toctree:: :caption: Project Links: :hidden: diff --git a/docs/source/about.rst b/docs/source/about.rst index d12a59de..993dec0c 100644 --- a/docs/source/about.rst +++ b/docs/source/about.rst @@ -7,408 +7,312 @@ About PrimAITE ============== +PrimAITE is a simulation environment for training agents to protect a computer network from cyber attacks. + Features ******** PrimAITE provides the following features: -* A flexible network / system laydown based on the Python networkx framework -* Nodes and links (edges) host Python classes in order to present attributes and methods (and hence, a more representative model of a platform / system) -* A 'green agent' Information Exchange Requirement (IER) function allows the representation of traffic (protocols and loading) on any / all links. Application of IERs is based on the status of node operating systems and services -* A 'green agent' node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. changing the Hardware state, Software State, Service state, or File System state) -* An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP, destination IP, protocol and port). Application of IERs adheres to any ACL restrictions -* Presents an OpenAI Gym interface to the environment, allowing integration with any OpenAI Gym compliant defensive agents -* Red agent activity based on 'red' IERs and 'red' PoL -* Defined reward function for use with RL agents (based on nodes status, and green / red IER success) -* Fully configurable (network / system laydown, IERs, node PoL, ACL, episode step period, episode max steps) and repeatable to suit the training requirements of agents. Therefore, not bound to a representation of any particular platform, system or technology -* Full capture of discrete metrics relating to agent training (full system state, agent actions taken, average reward) -* Networkx provides laydown visualisation capability - -Architecture - Nodes and Links -****************************** - -**Nodes** - -An inheritance model has been adopted in order to model nodes. All nodes have the following base attributes (Class: Node): - -* ID -* Name -* Type (e.g. computer, switch, RTU - enumeration) -* Priority (P1, P2, P3, P4 or P5 - enumeration) -* Hardware State (ON, OFF, RESETTING, SHUTTING_DOWN, BOOTING - enumeration) - -Active Nodes also have the following attributes (Class: Active Node): - -* IP Address -* Software State (GOOD, PATCHING, COMPROMISED - enumeration) -* File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration) - -Service Nodes also have the following attributes (Class: Service Node): - -* List of Services (where service is composed of service name and port). There is no theoretical limit on the number of services that can be modelled. Services and protocols are currently intrinsically linked (i.e. a service is an application on a node transmitting traffic of this protocol type) -* Service state (GOOD, PATCHING, COMPROMISED, OVERWHELMED - enumeration) - -Passive Nodes are currently not used (but may be employed for non IP-based components such as machinery actuators in future releases). - -**Links** - -Links are modelled both as network edges (networkx) and as Python classes, in order to extend their functionality. Links include the following attributes: - -* ID -* Name -* Bandwidth (bits/s) -* Source node ID -* Destination node ID -* Protocol list (containing the loading of protocols currently running on the link) - -When the simulation runs, IERs are applied to the links in order to model traffic loading, individually assigned to each protocol. This allows green (background) and red agent behaviour to be modelled, and defensive agents to identify suspicious traffic patterns at a protocol / traffic loading level of fidelity. - -Information Exchange Requirements (IERs) -**************************************** - -PrimAITE adopts the concept of Information Exchange Requirements (IERs) to model both green agent (background) and red agent (adversary) behaviour. IERs are used to initiate modelling of traffic loading on the network, and have the following attributes: - -* ID -* Start step (i.e. which step in the training episode should the IER start) -* End step (i.e. which step in the training episode should the IER end) -* Source node ID -* Destination node ID -* Load (bits/s) -* Protocol -* Port -* Running status (i.e. on / off) - -The application of green agent IERs between a source and destination follows a number of rules. Specifically: - -1. Does the current simulation time step fall between IER start and end step -2. Is the source node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING) -3. Is the destination node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING) -4. Are there any Access Control List rules in place that prevent the application of this IER -5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level) - -For red agent IERs, the application of IERs between a source and destination follows a number of subtly different rules. Specifically: - -1. Does the current simulation time step fall between IER start and end step -2. Is the source node operational, and is the service (protocol / port) associated with the IER (a) present on that node and (b) already in a compromised state -3. Is the destination node operational, and is the service (protocol / port) associated with the IER present on that node -4. Are there any Access Control List rules in place that prevent the application of this IER -5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level) - -Assuming the rules pass, the IER is applied to all relevant links (based on use of OSPF) between source and destination. - -Node Pattern-of-Life -******************** - -Every node can be impacted (i.e. have a status change applied to it) by either green agent pattern-of-life or red agent pattern-of-life. This is distinct from IERs, and allows for attacks (and defence) to be modelled purely within the confines of a node. - -The status changes that can be made to a node are as follows: - -* All Nodes: - - * Hardware State: - - * ON - * OFF - * RESETTING - when a status of resetting is entered, the node will automatically exit this state after a number of steps (as defined by the nodeResetDuration configuration item) after which it returns to an ON state - * BOOTING - * SHUTTING_DOWN - -* Active Nodes and Service Nodes: - - * Software State: - - * GOOD - * PATCHING - when a status of patching is entered, the node will automatically exit this state after a number of steps (as defined by the osPatchingDuration configuration item) after which it returns to a GOOD state - * COMPROMISED - - * File System State: - - * GOOD - * CORRUPT (can be resolved by repair or restore) - * DESTROYED (can be resolved by restore only) - * REPAIRING - when a status of repairing is entered, the node will automatically exit this state after a number of steps (as defined by the fileSystemRepairingLimit configuration item) after which it returns to a GOOD state - * RESTORING - when a status of repairing is entered, the node will automatically exit this state after a number of steps (as defined by the fileSystemRestoringLimit configuration item) after which it returns to a GOOD state - -* Service Nodes only: - - * Service State (for any associated service): - - * GOOD - * PATCHING - when a status of patching is entered, the service will automatically exit this state after a number of steps (as defined by the servicePatchingDuration configuration item) after which it returns to a GOOD state - * COMPROMISED - * OVERWHELMED - -Red agent pattern-of-life has an additional feature not found in the green pattern-of-life. This is the ability to influence the state of the attributes of a node via a number of different conditions: - - * DIRECT: - - The pattern-of-life described by the configuration file item will be applied regardless of any other conditions in the network. This is particularly useful for direct red agent entry into the network. - - * IER: - - The pattern-of-life described by the configuration file item will be applied to the service on the node, only if there is an IER of the same protocol / service type incoming at the specified timestep. - - * SERVICE: - - The pattern-of-life described by the configuration file item will be applied to the node based on the state of a service. The service can either be on the same node, or a different node within the network. - -Access Control List modelling -***************************** - -An Access Control List (ACL) is modelled to provide the means to manage traffic flows in the system. This will allow defensive agents the means to turn on / off rules, or potentially create new rules, to counter an attack. - -The ACL follows a standard network firewall format. For example: - -.. list-table:: ACL example - :widths: 25 25 25 25 25 - :header-rows: 1 - - * - Permission - - Source IP - - Dest IP - - Protocol - - Port - * - DENY - - 192.168.1.2 - - 192.168.1.3 - - HTTPS - - 443 - * - ALLOW - - 192.168.1.4 - - ANY - - SMTP - - 25 - * - DENY - - ANY - - 192.168.1.5 - - ANY - - ANY - -All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or ALLOW for the same parameters will override an earlier entry. - -Observation Spaces -****************** -The observation space provides the blue agent with information about the current status of nodes and links. - -PrimAITE builds on top of Gym Spaces to create an observation space that is easily configurable for users. It's made up of components which are managed by the :py:class:`primaite.environment.observations.ObservationsHandler`. Each training scenario can define its own observation space, and the user can choose which information to inlude, and how it should be formatted. - -NodeLinkTable component ------------------------ -For example, the :py:class:`primaite.environment.observations.NodeLinkTable` component represents the status of nodes and links as a ``gym.spaces.Box`` with an example format shown below: - -An example observation space is provided below: - -.. list-table:: Observation Space example - :widths: 25 25 25 25 25 25 25 - :header-rows: 1 - - * - - - ID - - Hardware State - - Software State - - File System State - - Service / Protocol A - - Service / Protocol B - * - Node A - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - * - Node B - - 2 - - 1 - - 3 - - 1 - - 1 - - 1 - * - Node C - - 3 - - 2 - - 1 - - 1 - - 3 - - 2 - * - Link 1 - - 5 - - 0 - - 0 - - 0 - - 0 - - 10000 - * - Link 2 - - 6 - - 0 - - 0 - - 0 - - 0 - - 10000 - * - Link 3 - - 7 - - 0 - - 0 - - 0 - - 5000 - - 0 - -For the nodes, the following values are represented: - -.. code-block:: - - [ - ID - Hardware State (1=ON, 2=OFF, 3=RESETTING, 4=SHUTTING_DOWN, 5=BOOTING) - Operating System State (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - File System State (0=none, 1=GOOD, 2=CORRUPT, 3=DESTROYED, 4=REPAIRING, 5=RESTORING) - Service1/Protocol1 state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - Service2/Protocol2 state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - ] - -(Note that each service available in the network is provided as a column, although not all nodes may utilise all services) - -For the links, the following statuses are represented: - -.. code-block:: - - [ - ID - Hardware State (0=not applicable) - Operating System State (0=not applicable) - File System State (0=not applicable) - Service1/Protocol1 state (Traffic load from this protocol on this link) - Service2/Protocol2 state (Traffic load from this protocol on this link) - ] - -NodeStatus component ----------------------- -This is a MultiDiscrete observation space that can be though of as a one-dimensional vector of discrete states. -The example above would have the following structure: - -.. code-block:: - - [ - node1_info - node2_info - node3_info - ] - -Each ``node_info`` contains the following: - -.. code-block:: - - [ - hardware_state (0=none, 1=ON, 2=OFF, 3=RESETTING, 4=SHUTTING_DOWN, 5=BOOTING) - software_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - file_system_state (0=none, 1=GOOD, 2=CORRUPT, 3=DESTROYED, 4=REPAIRING, 5=RESTORING) - service1_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - service2_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) - ] - -In a network with three nodes and two services, the full observation space would have 15 elements. It can be written with ``gym`` notation to indicate the number of discrete options for each of the elements of the observation space. For example: - -.. code-block:: - - gym.spaces.MultiDiscrete([4,5,6,4,4,4,5,6,4,4,4,5,6,4,4]) - -.. note:: - NodeStatus observation component provides information only about nodes. Links are not considered. - -LinkTrafficLevels ------------------ -This component is a MultiDiscrete space showing the traffic flow levels on the links in the network, after applying a threshold to convert it from a continuous to a discrete value. -There are two configurable parameters: -* ``quantisation_levels`` determines how many discrete bins to use for converting the continuous traffic value to discrete (default is 5). -* ``combine_service_traffic`` determines whether to separately output traffic use for each network protocol or whether to combine them into an overall value for the link. (default is ``True``) - -For example, with default parameters and a network with three links, the structure of this component would be: - -.. code-block:: - - [ - link1_status - link2_status - link3_status - ] - -Each ``link_status`` is a number from 0-4 representing the network load in relation to bandwidth. - -.. code-block:: - - 0 = No traffic (0%) - 1 = low traffic (1%-33%) - 2 = medium traffic (33%-66%) - 3 = high traffic (66%-99%) - 4 = max traffic/ overwhelmed (100%) - -Using ``gym`` notation, the shape of the obs space is: ``gym.spaces.MultiDiscrete([5,5,5])``. - - -Action Spaces -************** - -The action space available to the blue agent comes in two types: - - 1. Node-based - 2. Access Control List - 3. Any (Agent can take both node-based and ACL-based actions) - -The choice of action space used during a training session is determined in the config_[name].yaml file. - -**Node-Based** - -The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym spaces.Discrete type, as follows: - - * Dictionary item {... ,1: [x1, x2, x3,x4] ...} - The placeholders inside the list under the key '1' mean the following: - - * [0, num nodes] - Node ID (0 = nothing, node ID) - * [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = SoftwareState, 3 = service state, 4 = file system state) - * [0, 3] - Action on property (0 = nothing, 1 = on / scan, 2 = off / repair, 3 = reset / patch / restore) - * [0, num services] - Resolves to service ID (0 = nothing, resolves to service) - -**Access Control List** - -The blue agent is able to influence the configuration of the Access Control List rule set (which implements a system-wide firewall). In this instance, the action space is an OpenAI spaces.Discrete type, as follows: - - * Dictionary item {... ,1: [x1, x2, x3, x4, x5, x6] ...} - The placeholders inside the list under the key '1' mean the following: - - * [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - * [0, 1] - Permission (0 = DENY, 1 = ALLOW) - * [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - * [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - * [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - * [0, num ports] - Port (0 = any, then 1 -> x resolving to port) - -**ANY** -The agent is able to carry out both **Node-Based** and **Access Control List** operations. - -This means the dictionary will contain key-value pairs in the format of BOTH Node-Based and Access Control List as seen above. - -Rewards -******* - -A reward value is presented back to the blue agent on the conclusion of every step. The reward value is calculated via two methods which combine to give the total value: - - 1. Node and service status - 2. IER status - -**Node and service status** - -On every step, the status of each node is compared against both a reference environment (simulating the situation if the red and blue agents had not impacted the environment) -and the before and after state of the environment. If the comparison against the reference environment shows no difference, then the score provided is "AllOK". If there is a -difference with respect to the reference environment, the before and after states are compared, and a score determined. See :ref:`config` for details of reward values. - -**IER status** - -On every step, the full IER set is examined to determine whether green and red agent IERs are being permitted to run. Any red agent IERs running incur a penalty; any green agent -IERs not permitted to run also incur a penalty. See :ref:`config` for details of reward values. - -Future Enhancements -******************* - -The PrimAITE project has an ambition to include the following enhancements in future releases: - -* Integration with a suitable standardised framework to allow multi-agent integration -* Integration with external threat emulation tools, either using off-line data, or integrating at runtime +* A flexible system for defining network layouts and host configurations +* Highly configurable network hosts, including definition of software, file system, and network interfaces, +* Realistic network traffic simulation, including address and sending packets via internet protocols like TCP, UDP, ICMP, etc. +* Routers with traffic routing and firewall capabilities +* Interfaces with ARCD GATE to allow training of agents +* Simulation of customisable deterministic agents +* Support for multiple agents, each having their own customisable observation space, action space, and reward function definition. + + +Structure +********* + +PrimAITE consists of a simulator and a 'game' layer that allows agents to interact with the simulator. The simulator is built in a modular way where each component such as network hosts, links, networking devices, softwares, etc. are implemented as instances of a base class, meaning they all support the same interface. This allows for standardised configuration using either the Python API or YAML files. +The game layer is built on top of the simulator and it consumes the simulation action/state interface to allow agents to interact with the simulator. The game layer is also responsible for defining the reward function and observation space for the agents. + + +.. + Architecture - Nodes and Links + ****************************** + **Nodes** + An inheritance model has been adopted in order to model nodes. All nodes have the following base attributes (Class: Node): + * ID + * Name + * Type (e.g. computer, switch, RTU - enumeration) + * Priority (P1, P2, P3, P4 or P5 - enumeration) + * Hardware State (ON, OFF, RESETTING, SHUTTING_DOWN, BOOTING - enumeration) + Active Nodes also have the following attributes (Class: Active Node): + * IP Address + * Software State (GOOD, PATCHING, COMPROMISED - enumeration) + * File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration) + Service Nodes also have the following attributes (Class: Service Node): + * List of Services (where service is composed of service name and port). There is no theoretical limit on the number of services that can be modelled. Services and protocols are currently intrinsically linked (i.e. a service is an application on a node transmitting traffic of this protocol type) + * Service state (GOOD, PATCHING, COMPROMISED, OVERWHELMED - enumeration) + Passive Nodes are currently not used (but may be employed for non IP-based components such as machinery actuators in future releases). + **Links** + Links are modelled both as network edges (networkx) and as Python classes, in order to extend their functionality. Links include the following attributes: + * ID + * Name + * Bandwidth (bits/s) + * Source node ID + * Destination node ID + * Protocol list (containing the loading of protocols currently running on the link) + When the simulation runs, IERs are applied to the links in order to model traffic loading, individually assigned to each protocol. This allows green (background) and red agent behaviour to be modelled, and defensive agents to identify suspicious traffic patterns at a protocol / traffic loading level of fidelity. + Information Exchange Requirements (IERs) + **************************************** + PrimAITE adopts the concept of Information Exchange Requirements (IERs) to model both green agent (background) and red agent (adversary) behaviour. IERs are used to initiate modelling of traffic loading on the network, and have the following attributes: + * ID + * Start step (i.e. which step in the training episode should the IER start) + * End step (i.e. which step in the training episode should the IER end) + * Source node ID + * Destination node ID + * Load (bits/s) + * Protocol + * Port + * Running status (i.e. on / off) + The application of green agent IERs between a source and destination follows a number of rules. Specifically: + 1. Does the current simulation time step fall between IER start and end step + 2. Is the source node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING) + 3. Is the destination node operational (both physically and at an O/S level), and is the service (protocol / port) associated with the IER (a) present on this node, and (b) in an operational state (i.e. not PATCHING) + 4. Are there any Access Control List rules in place that prevent the application of this IER + 5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level) + For red agent IERs, the application of IERs between a source and destination follows a number of subtly different rules. Specifically: + 1. Does the current simulation time step fall between IER start and end step + 2. Is the source node operational, and is the service (protocol / port) associated with the IER (a) present on that node and (b) already in a compromised state + 3. Is the destination node operational, and is the service (protocol / port) associated with the IER present on that node + 4. Are there any Access Control List rules in place that prevent the application of this IER + 5. Are all switches in the (OSPF) path between source and destination operational (both physically and at an O/S level) + Assuming the rules pass, the IER is applied to all relevant links (based on use of OSPF) between source and destination. + Node Pattern-of-Life + ******************** + Every node can be impacted (i.e. have a status change applied to it) by either green agent pattern-of-life or red agent pattern-of-life. This is distinct from IERs, and allows for attacks (and defence) to be modelled purely within the confines of a node. + The status changes that can be made to a node are as follows: + * All Nodes: + * Hardware State: + * ON + * OFF + * RESETTING - when a status of resetting is entered, the node will automatically exit this state after a number of steps (as defined by the nodeResetDuration configuration item) after which it returns to an ON state + * BOOTING + * SHUTTING_DOWN + * Active Nodes and Service Nodes: + * Software State: + * GOOD + * PATCHING - when a status of patching is entered, the node will automatically exit this state after a number of steps (as defined by the osPatchingDuration configuration item) after which it returns to a GOOD state + * COMPROMISED + * File System State: + * GOOD + * CORRUPT (can be resolved by repair or restore) + * DESTROYED (can be resolved by restore only) + * REPAIRING - when a status of repairing is entered, the node will automatically exit this state after a number of steps (as defined by the fileSystemRepairingLimit configuration item) after which it returns to a GOOD state + * RESTORING - when a status of repairing is entered, the node will automatically exit this state after a number of steps (as defined by the fileSystemRestoringLimit configuration item) after which it returns to a GOOD state + * Service Nodes only: + * Service State (for any associated service): + * GOOD + * PATCHING - when a status of patching is entered, the service will automatically exit this state after a number of steps (as defined by the servicePatchingDuration configuration item) after which it returns to a GOOD state + * COMPROMISED + * OVERWHELMED + Red agent pattern-of-life has an additional feature not found in the green pattern-of-life. This is the ability to influence the state of the attributes of a node via a number of different conditions: + * DIRECT: + The pattern-of-life described by the configuration file item will be applied regardless of any other conditions in the network. This is particularly useful for direct red agent entry into the network. + * IER: + The pattern-of-life described by the configuration file item will be applied to the service on the node, only if there is an IER of the same protocol / service type incoming at the specified timestep. + * SERVICE: + The pattern-of-life described by the configuration file item will be applied to the node based on the state of a service. The service can either be on the same node, or a different node within the network. + Access Control List modelling + ***************************** + An Access Control List (ACL) is modelled to provide the means to manage traffic flows in the system. This will allow defensive agents the means to turn on / off rules, or potentially create new rules, to counter an attack. + The ACL follows a standard network firewall format. For example: + .. list-table:: ACL example + :widths: 25 25 25 25 25 + :header-rows: 1 + * - Permission + - Source IP + - Dest IP + - Protocol + - Port + * - DENY + - 192.168.1.2 + - 192.168.1.3 + - HTTPS + - 443 + * - ALLOW + - 192.168.1.4 + - ANY + - SMTP + - 25 + * - DENY + - ANY + - 192.168.1.5 + - ANY + - ANY + All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or ALLOW for the same parameters will override an earlier entry. + Observation Spaces + ****************** + The observation space provides the blue agent with information about the current status of nodes and links. + PrimAITE builds on top of Gym Spaces to create an observation space that is easily configurable for users. It's made up of components which are managed by the :py:class:`primaite.environment.observations.ObservationsHandler`. Each training scenario can define its own observation space, and the user can choose which information to inlude, and how it should be formatted. + NodeLinkTable component + ----------------------- + For example, the :py:class:`primaite.environment.observations.NodeLinkTable` component represents the status of nodes and links as a ``gym.spaces.Box`` with an example format shown below: + An example observation space is provided below: + .. list-table:: Observation Space example + :widths: 25 25 25 25 25 25 25 + :header-rows: 1 + * - + - ID + - Hardware State + - Software State + - File System State + - Service / Protocol A + - Service / Protocol B + * - Node A + - 1 + - 1 + - 1 + - 1 + - 1 + - 1 + * - Node B + - 2 + - 1 + - 3 + - 1 + - 1 + - 1 + * - Node C + - 3 + - 2 + - 1 + - 1 + - 3 + - 2 + * - Link 1 + - 5 + - 0 + - 0 + - 0 + - 0 + - 10000 + * - Link 2 + - 6 + - 0 + - 0 + - 0 + - 0 + - 10000 + * - Link 3 + - 7 + - 0 + - 0 + - 0 + - 5000 + - 0 + For the nodes, the following values are represented: + .. code-block:: + [ + ID + Hardware State (1=ON, 2=OFF, 3=RESETTING, 4=SHUTTING_DOWN, 5=BOOTING) + Operating System State (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + File System State (0=none, 1=GOOD, 2=CORRUPT, 3=DESTROYED, 4=REPAIRING, 5=RESTORING) + Service1/Protocol1 state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + Service2/Protocol2 state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + ] + (Note that each service available in the network is provided as a column, although not all nodes may utilise all services) + For the links, the following statuses are represented: + .. code-block:: + [ + ID + Hardware State (0=not applicable) + Operating System State (0=not applicable) + File System State (0=not applicable) + Service1/Protocol1 state (Traffic load from this protocol on this link) + Service2/Protocol2 state (Traffic load from this protocol on this link) + ] + NodeStatus component + ---------------------- + This is a MultiDiscrete observation space that can be though of as a one-dimensional vector of discrete states. + The example above would have the following structure: + .. code-block:: + [ + node1_info + node2_info + node3_info + ] + Each ``node_info`` contains the following: + .. code-block:: + [ + hardware_state (0=none, 1=ON, 2=OFF, 3=RESETTING, 4=SHUTTING_DOWN, 5=BOOTING) + software_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + file_system_state (0=none, 1=GOOD, 2=CORRUPT, 3=DESTROYED, 4=REPAIRING, 5=RESTORING) + service1_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + service2_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) + ] + In a network with three nodes and two services, the full observation space would have 15 elements. It can be written with ``gym`` notation to indicate the number of discrete options for each of the elements of the observation space. For example: + .. code-block:: + gym.spaces.MultiDiscrete([4,5,6,4,4,4,5,6,4,4,4,5,6,4,4]) + .. note:: + NodeStatus observation component provides information only about nodes. Links are not considered. + LinkTrafficLevels + ----------------- + This component is a MultiDiscrete space showing the traffic flow levels on the links in the network, after applying a threshold to convert it from a continuous to a discrete value. + There are two configurable parameters: + * ``quantisation_levels`` determines how many discrete bins to use for converting the continuous traffic value to discrete (default is 5). + * ``combine_service_traffic`` determines whether to separately output traffic use for each network protocol or whether to combine them into an overall value for the link. (default is ``True``) + For example, with default parameters and a network with three links, the structure of this component would be: + .. code-block:: + [ + link1_status + link2_status + link3_status + ] + Each ``link_status`` is a number from 0-4 representing the network load in relation to bandwidth. + .. code-block:: + 0 = No traffic (0%) + 1 = low traffic (1%-33%) + 2 = medium traffic (33%-66%) + 3 = high traffic (66%-99%) + 4 = max traffic/ overwhelmed (100%) + Using ``gym`` notation, the shape of the obs space is: ``gym.spaces.MultiDiscrete([5,5,5])``. + Action Spaces + ************** + The action space available to the blue agent comes in two types: + 1. Node-based + 2. Access Control List + 3. Any (Agent can take both node-based and ACL-based actions) + The choice of action space used during a training session is determined in the config_[name].yaml file. + **Node-Based** + The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym spaces.Discrete type, as follows: + * Dictionary item {... ,1: [x1, x2, x3,x4] ...} + The placeholders inside the list under the key '1' mean the following: + * [0, num nodes] - Node ID (0 = nothing, node ID) + * [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = SoftwareState, 3 = service state, 4 = file system state) + * [0, 3] - Action on property (0 = nothing, 1 = on / scan, 2 = off / repair, 3 = reset / patch / restore) + * [0, num services] - Resolves to service ID (0 = nothing, resolves to service) + **Access Control List** + The blue agent is able to influence the configuration of the Access Control List rule set (which implements a system-wide firewall). In this instance, the action space is an OpenAI spaces.Discrete type, as follows: + * Dictionary item {... ,1: [x1, x2, x3, x4, x5, x6] ...} + The placeholders inside the list under the key '1' mean the following: + * [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) + * [0, 1] - Permission (0 = DENY, 1 = ALLOW) + * [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) + * [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) + * [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) + * [0, num ports] - Port (0 = any, then 1 -> x resolving to port) + **ANY** + The agent is able to carry out both **Node-Based** and **Access Control List** operations. + This means the dictionary will contain key-value pairs in the format of BOTH Node-Based and Access Control List as seen above. + Rewards + ******* + A reward value is presented back to the blue agent on the conclusion of every step. The reward value is calculated via two methods which combine to give the total value: + 1. Node and service status + 2. IER status + **Node and service status** + On every step, the status of each node is compared against both a reference environment (simulating the situation if the red and blue agents had not impacted the environment) + and the before and after state of the environment. If the comparison against the reference environment shows no difference, then the score provided is "AllOK". If there is a + difference with respect to the reference environment, the before and after states are compared, and a score determined. See :ref:`config` for details of reward values. + **IER status** + On every step, the full IER set is examined to determine whether green and red agent IERs are being permitted to run. Any red agent IERs running incur a penalty; any green agent + IERs not permitted to run also incur a penalty. See :ref:`config` for details of reward values. + Future Enhancements + ******************* + The PrimAITE project has an ambition to include the following enhancements in future releases: + * Integration with a suitable standardised framework to allow multi-agent integration + * Integration with external threat emulation tools, either using off-line data, or integrating at runtime diff --git a/docs/source/action_system.rst b/docs/source/action_system.rst deleted file mode 100644 index 88baf232..00000000 --- a/docs/source/action_system.rst +++ /dev/null @@ -1,88 +0,0 @@ -.. only:: comment - - © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -Actions System -============== - -``SimComponent``s in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``Action``. - -Just like other aspects of SimComponent, the actions are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved with the following design decisions: - -- API - An 'action' contains two elements: - - 1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '', 'service', '', 'restart']`. - 2. ``context`` - optional extra information that can be used to decide how to process the action. This is formatted as a dictionary. For example, if the action requires authentication, the context can include information about the user that initiated the request to decide if their permissions are sufficient. - -- request - The request is a list of strings which help specify who should handle the request. The strings in the request list help RequestManagers traverse the 'ownership tree' of SimComponent. The example given above would be handled in the following way: - - 1. ``Simulation`` receives `['network', 'node', '', 'service', '', 'restart']`. - The first element of the action is ``network``, therefore it passes the action down to its network. - 2. ``Network`` receives `['node', '', 'service', '', 'restart']`. - The first element of the action is ``node``, therefore the network looks at the node uuid and passes the action down to the node with that uuid. - 3. ``Node`` receives `['service', '', 'restart']`. - The first element of the action is ``service``, therefore the node looks at the service uuid and passes the rest of the action to the service with that uuid. - 4. ``Service`` receives ``['restart']``. - Since ``restart`` is a defined action in the service's own RequestManager, the service performs a restart. - -Technical Detail -================ - -This system was achieved by implementing two classes, :py:class:`primaite.simulator.core.Action`, and :py:class:`primaite.simulator.core.RequestManager`. - -Action ------- - -The ``Action`` object stores a reference to a method that performs the action, for example a node could have an action that stores a reference to ``self.turn_on()``. Technically, this can be any callable that accepts `request, context` as it's parameters. In practice, this is often defined using ``lambda`` functions within a component's ``self._init_request_manager()`` method. Optionally, the ``Action`` object can also hold a validator that will permit/deny the action depending on context. - -RequestManager -------------- - -The ``RequestManager`` object stores a mapping between strings and actions. It is responsible for processing the ``request`` and passing it down the ownership tree. Technically, the ``RequestManager`` is itself a callable that accepts `request, context` tuple, and so it can be chained with other action managers. - -A simple example without chaining can be seen in the :py:class:`primaite.simulator.file_system.file_system.File` class. - -.. code-block:: python - - class File(FileSystemItemABC): - ... - def _init_request_manager(self): - ... - request_manager.add_request("scan", Action(func=lambda request, context: self.scan())) - request_manager.add_request("repair", Action(func=lambda request, context: self.repair())) - request_manager.add_request("restore", Action(func=lambda request, context: self.restore())) - -*ellipses (``...``) used to omit code impertinent to this explanation* - -Chaining RequestManagers ------------------------ - -Since the method for performing an action needs to accept `request, context` as parameters, and RequestManager itself is a callable that accepts `request, context` as parameters, it possible to use RequestManager as an action. In fact, that is how PrimAITE deals with traversing the ownership tree. Each time an RequestManager accepts a request, it pops the first elements and uses it to decide to which Action it should send the remaining request. However, the Action could have another RequestManager as it's function, therefore the request will be routed again. Each time the request is passed to a new action manager, the first element is popped. - -An example of how this works is in the :py:class:`primaite.simulator.network.hardware.base.Node` class. - -.. code-block:: python - - class Node(SimComponent): - ... - def _init_request_manager(self): - ... - # a regular action which is processed by the Node itself - request_manager.add_request("turn_on", Action(func=lambda request, context: self.turn_on())) - - # if the Node receives a request where the first word is 'service', it will use a dummy manager - # called self._service_request_manager to pass on the reqeust to the relevant service. This dummy - # manager is simply here to map the service UUID that that service's own action manager. This is - # done because the next string after "service" is always the uuid of that service, so we need an - # RequestManager to pop that string before sending it onto the relevant service's RequestManager. - self._service_request_manager = RequestManager() - request_manager.add_request("service", Action(func=self._service_request_manager)) - ... - - def install_service(self, service): - self.services[service.uuid] = service - ... - # Here, the service UUID is registered to allow passing actions between the node and the service. - self._service_request_manager.add_request(service.uuid, Action(func=service._request_manager)) diff --git a/docs/source/config(v2).rst b/docs/source/config(v2).rst deleted file mode 100644 index 35233cf5..00000000 --- a/docs/source/config(v2).rst +++ /dev/null @@ -1,491 +0,0 @@ -.. only:: comment - - © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -.. _config: - -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** - - 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/custom_agent.rst b/docs/source/custom_agent.rst index 040b4b3d..0a08ae74 100644 --- a/docs/source/custom_agent.rst +++ b/docs/source/custom_agent.rst @@ -11,132 +11,4 @@ Integrating a user defined blue agent .. note:: - If you are planning to implement custom RL agents into PrimAITE, you must use the project as a repository. If you install PrimAITE as a python package from wheel, custom agents are not supported. - -PrimAITE has integration with Ray RLLib and StableBaselines3 agents. All agents interface with PrimAITE through an :py:class:`primaite.agents.agent_abc.AgentSessionABC` which provides Input/Output of agent savefiles, as well as capturing and plotting performance metrics during training and evaluation. If you wish to integrate a custom blue agent, it is recommended to create a subclass of the :py:class:`primaite.agents.agent_abc.AgentSessionABC` and implement the ``__init__()``, ``_setup()``, ``_save_checkpoint()``, ``learn()``, ``evaluate()``, ``_get_latest_checkpoint``, ``load()``, and ``save()`` methods. - -Below is a barebones example of a custom agent implementation: - -.. code:: python - - # src/primaite/agents/my_custom_agent.py - - from primaite.agents.agent_abc import AgentSessionABC - from primaite.common.enums import AgentFramework, AgentIdentifier - - class CustomAgent(AgentSessionABC): - def __init__(self, training_config_path, lay_down_config_path): - super().__init__(training_config_path, lay_down_config_path) - assert self._training_config.agent_framework == AgentFramework.CUSTOM - assert self._training_config.agent_identifier == AgentIdentifier.MY_AGENT - self._setup() - - def _setup(self): - super()._setup() - self._env = Primaite( - training_config_path=self._training_config_path, - lay_down_config_path=self._lay_down_config_path, - session_path=self.session_path, - timestamp_str=self.timestamp_str, - ) - self._agent = ... # your code to setup agent - - def _save_checkpoint(self): - checkpoint_num = self._training_config.checkpoint_every_n_episodes - episode_count = self._env.episode_count - save_checkpoint = False - if checkpoint_num: - save_checkpoint = episode_count % checkpoint_num == 0 - # saves checkpoint if the episode count is not 0 and save_checkpoint flag was set to true - if episode_count and save_checkpoint: - ... - # your code to save checkpoint goes here. - # The path should start with self.checkpoints_path and include the episode number. - - def learn(self): - ... - # call your agent's learning function here. - - super().learn() # this will finalise learning and output session metadata - self.save() - - def evaluate(self): - ... - # call your agent's evaluation function here. - - self._env.close() - super().evaluate() - - def _get_latest_checkpoint(self): - ... - # Load an agent from file. - - @classmethod - def load(cls, path): - ... - # Create a CustomAgent object which loads model weights from file. - - def save(self): - ... - # Call your agent's function that saves it to a file - - -You will also need to modify :py:class:`primaite.primaite_session.PrimaiteSession` and :py:mod:`primaite.common.enums` to capture your new agent identifiers. - -.. code-block:: python - :emphasize-lines: 17, 18 - - # src/primaite/common/enums.py - - class AgentIdentifier(Enum): - """The Red Agent algo/class.""" - A2C = 1 - "Advantage Actor Critic" - PPO = 2 - "Proximal Policy Optimization" - HARDCODED = 3 - "The Hardcoded agents" - DO_NOTHING = 4 - "The DoNothing agents" - RANDOM = 5 - "The RandomAgent" - DUMMY = 6 - "The DummyAgent" - CUSTOM_AGENT = 7 - "Your custom agent" - -.. code-block:: python - :emphasize-lines: 3, 11, 12 - - # src/primaite_session.py - - from primaite.agents.my_custom_agent import CustomAgent - - # ... - - def setup(self): - """Performs the session setup.""" - if self._training_config.agent_framework == AgentFramework.CUSTOM: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.CUSTOM}") - if self._training_config.agent_identifier == AgentIdentifier.CUSTOM_AGENT: - self._agent_session = CustomAgent(self._training_config_path, self._lay_down_config_path) - if self._training_config.agent_identifier == AgentIdentifier.HARDCODED: - _LOGGER.debug(f"PrimaiteSession Setup: Agent Identifier =" f" {AgentIdentifier.HARDCODED}") - if self._training_config.action_type == ActionType.NODE: - # Deterministic Hardcoded Agent with Node Action Space - self._agent_session = HardCodedNodeAgent(self._training_config_path, self._lay_down_config_path) - -Finally, specify your agent in your training config. - -.. code-block:: yaml - - # ~/primaite/2.0.0/config/path/to/your/config_main.yaml - - # Training Config File - - agent_framework: CUSTOM - agent_identifier: CUSTOM_AGENT - random_red_agent: False - # ... - -Now you can :ref:`run a primaite session` with your custom agent by passing in the custom ``config_main``. + PrimAITE uses ARCD GATE for agent integration. In order to use a custom agent with PrimAITE, you must integrate it with ARCD GATE. Please look at the ARCD GATE documentation for more information. diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 0801c79e..aebabf66 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -11,7 +11,7 @@ Getting Started Pre-Requisites -In order to get **PrimAITE** installed, you will need to have a python version between 3.8 and 3.10 installed. If you don't already have it, this is how to install it: +In order to get **PrimAITE** installed, you will need to have a python version between 3.8 and 3.11 installed. If you don't already have it, this is how to install it: .. code-block:: bash @@ -33,39 +33,36 @@ In order to get **PrimAITE** installed, you will need to have a python version b Install PrimAITE **************** -1. Create a primaite directory in your home directory: - - +1. Create a directory for your PrimAITE project: .. code-block:: bash :caption: Unix - mkdir ~/primaite/2.0.0 + mkdir ~/primaite/3.0.0 .. code-block:: powershell :caption: Windows (Powershell) - mkdir ~\primaite\2.0.0 + mkdir ~\primaite\3.0.0 + 2. Navigate to the primaite directory and create a new python virtual environment (venv) - - .. code-block:: bash :caption: Unix - cd ~/primaite/2.0.0 + cd ~/primaite/3.0.0 python3 -m venv .venv .. code-block:: powershell :caption: Windows (Powershell) - cd ~\primaite\2.0.0 + cd ~\primaite\3.0.0 python3 -m venv .venv attrib +h .venv /s /d # Hides the .venv directory -3. Activate the venv +3. Activate the venv .. code-block:: bash :caption: Unix @@ -78,21 +75,34 @@ Install PrimAITE .\.venv\Scripts\activate -4. Install PrimAITE using pip from PyPi +4. Install PrimAITE from your saved wheel file + +.. code-block:: bash + :caption: Unix + + pip install path/to/your/primaite.whl + +.. code-block:: powershell + :caption: Windows (Powershell) + + pip install path\to\your\primaite.whl + + +5. Install ARCD GATE from wheel file .. code-block:: bash :caption: Unix - pip install primaite + pip install path/to/your/arcd_gate-0.1.0-py3-none-any.whl .. code-block:: powershell :caption: Windows (Powershell) - pip install primaite + pip install path\to\your\arcd_gate-0.1.0-py3-none-any.whl -5. Perform the PrimAITE setup +6. Perform the PrimAITE setup .. code-block:: bash :caption: Unix @@ -110,13 +120,14 @@ Clone & Install PrimAITE for Development To be able to extend PrimAITE further, or to build wheels manually before install, clone the repository to a location of your choice: +1. Clone the repository + .. code-block:: bash git clone https://github.com/Autonomous-Resilient-Cyber-Defence/PrimAITE cd primaite -Create and activate your Python virtual environment (venv) - +2. Create and activate your Python virtual environment (venv) .. code-block:: bash :caption: Unix @@ -130,8 +141,7 @@ Create and activate your Python virtual environment (venv) python3 -m venv venv .\venv\Scripts\activate -Install PrimAITE with the dev extra - +3. Install PrimAITE with the dev extra .. code-block:: bash :caption: Unix @@ -144,4 +154,16 @@ Install PrimAITE with the dev extra pip install -e .[dev] +4. Install ARCD GATE from wheel file + +.. code-block:: bash + :caption: Unix + + pip install GATE/arcd_gate-0.1.0-py3-none-any.whl + +.. code-block:: powershell + :caption: Windows (Powershell) + + pip install GATE\arcd_gate-0.1.0-py3-none-any.whl + To view the complete list of packages installed during PrimAITE installation, go to the dependencies page (:ref:`Dependencies`). diff --git a/docs/source/migration_1.2_-_2.0.rst b/docs/source/migration_1.2_-_2.0.rst deleted file mode 100644 index c38fcbe9..00000000 --- a/docs/source/migration_1.2_-_2.0.rst +++ /dev/null @@ -1,57 +0,0 @@ -.. only:: comment - - © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -v1.2 to v2.0 Migration guide -============================ - -**1. Installing PrimAITE** - - Like before, you can install primaite from the repository by running ``pip install -e .``. But, there is now an additional setup step which does several things, like setting up user directories, copy default configs and notebooks, etc. Once you have installed PrimAITE to your virtual environment, run this command to finalise setup. - - .. code-block:: bash - - primaite setup - -**2. Running a training session** - - In version 1.2 of PrimAITE, the main entry point for training or evaluating agents was the ``src/primaite/main.py`` file. v2.0.0 introduced managed 'sessions' which are responsible for reading configuration files, performing training, and writing outputs. - - ``main.py`` file still runs a training session but it now uses the new `PrimaiteSession`, and it now requires you to provide the path to your config files. - - .. code-block:: bash - - python src/primaite/main.py --tc path/to/training-config.yaml --ldc path/to/laydown-config.yaml - - Alternatively, the session can be invoked via the commandline by running: - - .. code-block:: bash - - primaite session --tc path/to/training-config.yaml --ldc path/to/laydown-config.yaml - -**3. Location of configs** - - In version 1.2, training configs and laydown configs were all stored in the project repository under ``src/primaite/config``. Version 2.0.0 introduced user data directories, and now when you install and setup PrimAITE, config files are stored in your user data location. On Linux/OSX, this is stored in ``~/primaite/2.0.0/config``. On Windows, this is stored in ``C:\Users\\primaite\configs``. Upon first setup, the configs folder is populated with some default yaml files. It is recommended that you store all your custom configuration files here. - -**4. Contents of configs** - - Some things that were previously part of the laydown config are now part of the traning config. - - * Actions - - If you have custom configs which use these, you will need to adapt them by moving the configuration from the laydown config to the training config. - - Also, there are new configurable items in the training config: - - * Observations - * Agent framework - * Agent - * Deep learning framework - * random red agents - * seed - * deterministic - * hard coded agent view - - Each of these items have default values which are designed so that PrimAITE has the same behaviour as it did in 1.2.0, so you do not have to specify them. - - ACL Rules in laydown configs have a new required parameter: ``position``. The lower the position, the higher up in the ACL table the rule will placed. If you have custom laydowns, you will need to go through them and add a position to each ACL_RULE. diff --git a/docs/source/primaite_session.rst b/docs/source/primaite_session.rst index 8ccc9070..472a361f 100644 --- a/docs/source/primaite_session.rst +++ b/docs/source/primaite_session.rst @@ -14,199 +14,200 @@ A PrimAITE session can be ran either with the ``primaite session`` command from (See :func:`primaite.cli.session`), or by calling :func:`primaite.main.run` from a Python terminal or Jupyter Notebook. Both the ``primaite session`` and :func:`primaite.main.run` take a training config and a lay down config as parameters. +.. note:: + 🚧 *UNDER CONSTRUCTION* 🚧 - - -.. code-block:: bash - :caption: Unix CLI - - cd ~/primaite/2.0.0 - source ./.venv/bin/activate - primaite session --tc ./config/my_training_config.yaml --ldc ./config/my_lay_down_config.yaml - -.. code-block:: powershell - :caption: Powershell CLI - - cd ~\primaite\2.0.0 - .\.venv\Scripts\activate - primaite session --tc .\config\my_training_config.yaml --ldc .\config\my_lay_down_config.yaml - - -.. code-block:: python - :caption: Python - - from primaite.main import run - - training_config = - lay_down_config = - run(training_config, lay_down_config) - -When a session is ran, a session output sub-directory is created in the users app sessions directory (``~/primaite/2.0.0/sessions``). -The sub-directory is formatted as such: ``~/primaite/2.0.0/sessions//_/`` - -For example, when running a session at 17:30:00 on 31st January 2023, the session will output to: -``~/primaite/2.0.0/sessions/2023-01-31/2023-01-31_17-30-00/``. - -``primaite session`` can be ran in the terminal/command prompt without arguments. It will use the default configs in the directory ``primaite/config/example_config``. - -To run a PrimAITE session using legacy training or laydown config files, add the ``--legacy-tc`` and/or ``legacy-ldc`` options. - - - -.. code-block:: bash - :caption: Unix CLI - - cd ~/primaite/2.0.0 - source ./.venv/bin/activate - primaite session --tc ./config/my_legacy_training_config.yaml --legacy-tc --ldc ./config/my_legacy_lay_down_config.yaml --legacy-ldc - -.. code-block:: powershell - :caption: Powershell CLI - - cd ~\primaite\2.0.0 - .\.venv\Scripts\activate - primaite session --tc .\config\my_legacy_training_config.yaml --legacy-tc --ldc .\config\my_legacy_lay_down_config.yaml --legacy-ldc - - -.. code-block:: python - :caption: Python - - from primaite.main import run - - training_config = - lay_down_config = - run(training_config, lay_down_config, legacy_training_config=True, legacy_lay_down_config=True) - - - - -Outputs -------- - -PrimAITE produces four types of outputs: - -* Session Metadata -* Results -* Diagrams -* Saved agents (training checkpoints and a final trained agent) - - -**Session Metadata** - -PrimAITE creates a ``session_metadata.json`` file that contains the following metadata: - - * **uuid** - The UUID assigned to the session upon instantiation. - * **start_datetime** - The date & time the session started in iso format. - * **end_datetime** - The date & time the session ended in iso format. - * **learning** - * **total_episodes** - The total number of training episodes completed. - * **total_time_steps** - The total number of training time steps completed. - * **evaluation** - * **total_episodes** - The total number of evaluation episodes completed. - * **total_time_steps** - The total number of evaluation time steps completed. - * **env** - * **training_config** - * **All training config items** - * **lay_down_config** - * **All lay down config items** - - -**Results** - -PrimAITE automatically creates two sets of results from each learning and evaluation session: - -* Average reward per episode - a csv file listing the average reward for each episode of the session. This provides, for example, an indication of the change over a training session of the reward value -* All transactions - a csv file listing the following values for every step of every episode: - - * Timestamp - * Episode number - * Step number - * Reward value - * Action taken (as presented by the blue agent on this step). Individual elements of the action space are presented in the format AS_X - * Initial observation space (what the blue agent observed when it decided its action) - -**Diagrams** - -* For each session, PrimAITE automatically creates a visualisation of the system / network lay down configuration. -* For each learning and evaluation task within the session, PrimAITE automatically plots the average reward per episode using PlotLY and saves it to the learning or evaluation subdirectory in the session directory. - -**Saved agents** - -For each training session, assuming the agent being trained implements the *save()* function and this function is called by the code, PrimAITE automatically saves the agent state. - -**Example Session Directory Structure** - -.. code-block:: text - - ~/ - └── primaite/ - └── 2.0.0/ - └── sessions/ - └── 2023-07-18/ - └── 2023-07-18_11-06-04/ - ├── evaluation/ - │ ├── all_transactions_2023-07-18_11-06-04.csv - │ ├── average_reward_per_episode_2023-07-18_11-06-04.csv - │ └── average_reward_per_episode_2023-07-18_11-06-04.png - ├── learning/ - │ ├── all_transactions_2023-07-18_11-06-04.csv - │ ├── average_reward_per_episode_2023-07-18_11-06-04.csv - │ ├── average_reward_per_episode_2023-07-18_11-06-04.png - │ ├── checkpoints/ - │ │ └── sb3ppo_10.zip - │ ├── SB3_PPO.zip - │ └── tensorboard_logs/ - │ ├── PPO_1/ - │ │ └── events.out.tfevents.1689674765.METD-9PMRFB3.42960.0 - │ ├── PPO_2/ - │ │ └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.1 - │ ├── PPO_3/ - │ │ └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.2 - │ ├── PPO_4/ - │ │ └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.3 - │ ├── PPO_5/ - │ │ └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.4 - │ ├── PPO_6/ - │ │ └── events.out.tfevents.1689674768.METD-9PMRFB3.42960.5 - │ ├── PPO_7/ - │ │ └── events.out.tfevents.1689674768.METD-9PMRFB3.42960.6 - │ ├── PPO_8/ - │ │ └── events.out.tfevents.1689674769.METD-9PMRFB3.42960.7 - │ ├── PPO_9/ - │ │ └── events.out.tfevents.1689674770.METD-9PMRFB3.42960.8 - │ └── PPO_10/ - │ └── events.out.tfevents.1689674770.METD-9PMRFB3.42960.9 - ├── network_2023-07-18_11-06-04.png - └── session_metadata.json - -Loading a session ------------------ - -A previous session can be loaded by providing the **directory** of the previous session to either the ``primaite session`` command from the cli -(See :func:`primaite.cli.session`), or by calling :func:`primaite.main.run` with session_path. - -.. tabs:: - - .. code-tab:: bash +.. + .. code-block:: bash :caption: Unix CLI cd ~/primaite/2.0.0 source ./.venv/bin/activate - primaite session --load "path/to/session" + primaite session --tc ./config/my_training_config.yaml --ldc ./config/my_lay_down_config.yaml - .. code-tab:: bash + .. code-block:: powershell :caption: Powershell CLI cd ~\primaite\2.0.0 .\.venv\Scripts\activate - primaite session --load "path\to\session" + primaite session --tc .\config\my_training_config.yaml --ldc .\config\my_lay_down_config.yaml - .. code-tab:: python + .. code-block:: python :caption: Python from primaite.main import run - run(session_path=) + training_config = + lay_down_config = + run(training_config, lay_down_config) -When PrimAITE runs a loaded session, PrimAITE will output in the provided session directory + When a session is ran, a session output sub-directory is created in the users app sessions directory (``~/primaite/2.0.0/sessions``). + The sub-directory is formatted as such: ``~/primaite/2.0.0/sessions//_/`` + + For example, when running a session at 17:30:00 on 31st January 2023, the session will output to: + ``~/primaite/2.0.0/sessions/2023-01-31/2023-01-31_17-30-00/``. + + ``primaite session`` can be ran in the terminal/command prompt without arguments. It will use the default configs in the directory ``primaite/config/example_config``. + + To run a PrimAITE session using legacy training or laydown config files, add the ``--legacy-tc`` and/or ``legacy-ldc`` options. + + + + .. code-block:: bash + :caption: Unix CLI + + cd ~/primaite/2.0.0 + source ./.venv/bin/activate + primaite session --tc ./config/my_legacy_training_config.yaml --legacy-tc --ldc ./config/my_legacy_lay_down_config.yaml --legacy-ldc + + .. code-block:: powershell + :caption: Powershell CLI + + cd ~\primaite\2.0.0 + .\.venv\Scripts\activate + primaite session --tc .\config\my_legacy_training_config.yaml --legacy-tc --ldc .\config\my_legacy_lay_down_config.yaml --legacy-ldc + + + .. code-block:: python + :caption: Python + + from primaite.main import run + + training_config = + lay_down_config = + run(training_config, lay_down_config, legacy_training_config=True, legacy_lay_down_config=True) + + + + + Outputs + ------- + + PrimAITE produces four types of outputs: + + * Session Metadata + * Results + * Diagrams + * Saved agents (training checkpoints and a final trained agent) + + + **Session Metadata** + + PrimAITE creates a ``session_metadata.json`` file that contains the following metadata: + + * **uuid** - The UUID assigned to the session upon instantiation. + * **start_datetime** - The date & time the session started in iso format. + * **end_datetime** - The date & time the session ended in iso format. + * **learning** + * **total_episodes** - The total number of training episodes completed. + * **total_time_steps** - The total number of training time steps completed. + * **evaluation** + * **total_episodes** - The total number of evaluation episodes completed. + * **total_time_steps** - The total number of evaluation time steps completed. + * **env** + * **training_config** + * **All training config items** + * **lay_down_config** + * **All lay down config items** + + + **Results** + + PrimAITE automatically creates two sets of results from each learning and evaluation session: + + * Average reward per episode - a csv file listing the average reward for each episode of the session. This provides, for example, an indication of the change over a training session of the reward value + * All transactions - a csv file listing the following values for every step of every episode: + + * Timestamp + * Episode number + * Step number + * Reward value + * Action taken (as presented by the blue agent on this step). Individual elements of the action space are presented in the format AS_X + * Initial observation space (what the blue agent observed when it decided its action) + + **Diagrams** + + * For each session, PrimAITE automatically creates a visualisation of the system / network lay down configuration. + * For each learning and evaluation task within the session, PrimAITE automatically plots the average reward per episode using PlotLY and saves it to the learning or evaluation subdirectory in the session directory. + + **Saved agents** + + For each training session, assuming the agent being trained implements the *save()* function and this function is called by the code, PrimAITE automatically saves the agent state. + + **Example Session Directory Structure** + + .. code-block:: text + + ~/ + └── primaite/ + └── 2.0.0/ + └── sessions/ + └── 2023-07-18/ + └── 2023-07-18_11-06-04/ + ├── evaluation/ + │ ├── all_transactions_2023-07-18_11-06-04.csv + │ ├── average_reward_per_episode_2023-07-18_11-06-04.csv + │ └── average_reward_per_episode_2023-07-18_11-06-04.png + ├── learning/ + │ ├── all_transactions_2023-07-18_11-06-04.csv + │ ├── average_reward_per_episode_2023-07-18_11-06-04.csv + │ ├── average_reward_per_episode_2023-07-18_11-06-04.png + │ ├── checkpoints/ + │ │ └── sb3ppo_10.zip + │ ├── SB3_PPO.zip + │ └── tensorboard_logs/ + │ ├── PPO_1/ + │ │ └── events.out.tfevents.1689674765.METD-9PMRFB3.42960.0 + │ ├── PPO_2/ + │ │ └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.1 + │ ├── PPO_3/ + │ │ └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.2 + │ ├── PPO_4/ + │ │ └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.3 + │ ├── PPO_5/ + │ │ └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.4 + │ ├── PPO_6/ + │ │ └── events.out.tfevents.1689674768.METD-9PMRFB3.42960.5 + │ ├── PPO_7/ + │ │ └── events.out.tfevents.1689674768.METD-9PMRFB3.42960.6 + │ ├── PPO_8/ + │ │ └── events.out.tfevents.1689674769.METD-9PMRFB3.42960.7 + │ ├── PPO_9/ + │ │ └── events.out.tfevents.1689674770.METD-9PMRFB3.42960.8 + │ └── PPO_10/ + │ └── events.out.tfevents.1689674770.METD-9PMRFB3.42960.9 + ├── network_2023-07-18_11-06-04.png + └── session_metadata.json + + Loading a session + ----------------- + + A previous session can be loaded by providing the **directory** of the previous session to either the ``primaite session`` command from the cli + (See :func:`primaite.cli.session`), or by calling :func:`primaite.main.run` with session_path. + + .. tabs:: + + .. code-tab:: bash + :caption: Unix CLI + + cd ~/primaite/2.0.0 + source ./.venv/bin/activate + primaite session --load "path/to/session" + + .. code-tab:: bash + :caption: Powershell CLI + + cd ~\primaite\2.0.0 + .\.venv\Scripts\activate + primaite session --load "path\to\session" + + + .. code-tab:: python + :caption: Python + + from primaite.main import run + + run(session_path=) + + When PrimAITE runs a loaded session, PrimAITE will output in the provided session directory diff --git a/docs/source/request_system.rst b/docs/source/request_system.rst new file mode 100644 index 00000000..41d8eec4 --- /dev/null +++ b/docs/source/request_system.rst @@ -0,0 +1,90 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +Request System +============== + +``SimComponent``s in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``. + +Just like other aspects of SimComponent, the request typess are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved in the following way: + +- API + An ``RequestType`` contains two elements: + + 1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '', 'service', '', 'restart']`. + 2. ``context`` - optional extra information that can be used to decide how to process the request. This is formatted as a dictionary. For example, if the request requires authentication, the context can include information about the user that initiated the request to decide if their permissions are sufficient. + +- request + The request is a list of strings which help specify who should handle the request. The strings in the request list help RequestManagers traverse the 'ownership tree' of SimComponent. The example given above would be handled in the following way: + + 1. ``Simulation`` receives `['network', 'node', '', 'service', '', 'restart']`. + The first element of the request is ``network``, therefore it passes the request down to its network. + 2. ``Network`` receives `['node', '', 'service', '', 'restart']`. + The first element of the request is ``node``, therefore the network looks at the node uuid and passes the request down to the node with that uuid. + 3. ``Node`` receives `['service', '', 'restart']`. + The first element of the request is ``service``, therefore the node looks at the service uuid and passes the rest of the request to the service with that uuid. + 4. ``Service`` receives ``['restart']``. + Since ``restart`` is a defined request type in the service's own RequestManager, the service performs a restart. + +Technical Detail +================ + +This system was achieved by implementing two classes, :py:class:`primaite.simulator.core.RequestType`, and :py:class:`primaite.simulator.core.RequestManager`. + +``RequestType`` +------ + +The ``RequestType`` object stores a reference to a method that executes the request, for example a node could have a request type that stores a reference to ``self.turn_on()``. Technically, this can be any callable that accepts `request, context` as it's parameters. In practice, this is often defined using ``lambda`` functions within a component's ``self._init_request_manager()`` method. Optionally, the ``RequestType`` object can also hold a validator that will permit/deny the request depending on context. + +``RequestManager`` +------------- + +The ``RequestManager`` object stores a mapping between strings and request types. It is responsible for processing the request and passing it down the ownership tree. Technically, the ``RequestManager`` is itself a callable that accepts `request, context` tuple, and so it can be chained with other request managers. + +A simple example without chaining can be seen in the :py:class:`primaite.simulator.file_system.file_system.File` class. + +.. code-block:: python + + class File(FileSystemItemABC): + ... + def _init_request_manager(self): + ... + request_manager.add_request("scan", RequestType(func=lambda request, context: self.scan())) + request_manager.add_request("repair", RequestType(func=lambda request, context: self.repair())) + request_manager.add_request("restore", RequestType(func=lambda request, context: self.restore())) + +*ellipses (``...``) used to omit code impertinent to this explanation* + +Chaining RequestManagers +----------------------- + +A request function needs to be a callable that accepts ``request, context`` as parameters. Since the request manager resolves requests by invoking it with ``request, context`` as parameter, it is possible to use a ``RequestManager`` as a ``RequestType``. + +When a RequestManager accepts a request, it pops the first element and uses it to decide where it should send the remaining request. This is how PrimAITE traverses the ownership tree. If the ``RequestType`` has another ``RequestManager`` as its function, the request will be routed again. Each time the request is passed to a new request manager, the first element is popped. + +An example of how this works is in the :py:class:`primaite.simulator.network.hardware.base.Node` class. + +.. code-block:: python + + class Node(SimComponent): + ... + def _init_request_manager(self): + ... + # a regular action which is processed by the Node itself + request_manager.add_request("turn_on", RequestType(func=lambda request, context: self.turn_on())) + + # if the Node receives a request where the first word is 'service', it will use a dummy manager + # called self._service_request_manager to pass on the reqeust to the relevant service. This dummy + # manager is simply here to map the service UUID that that service's own action manager. This is + # done because the next string after "service" is always the uuid of that service, so we need an + # RequestManager to pop that string before sending it onto the relevant service's RequestManager. + self._service_request_manager = RequestManager() + request_manager.add_request("service", RequestType(func=self._service_request_manager)) + ... + + def install_service(self, service): + self.services[service.uuid] = service + ... + # Here, the service UUID is registered to allow passing actions between the node and the service. + self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager)) diff --git a/docs/source/simulation.rst b/docs/source/simulation.rst index 8671a2d2..5e259c6f 100644 --- a/docs/source/simulation.rst +++ b/docs/source/simulation.rst @@ -23,4 +23,8 @@ Contents simulation_components/network/network simulation_components/system/internal_frame_processing simulation_components/system/software - action_system + simulation_components/system/data_manipulation_bot + simulation_components/system/database_client_server + simulation_components/system/dns_client_server + simulation_components/system/ftp_client_server + simulation_components/system/web_browser_and_web_server_service diff --git a/docs/source/simulation_structure.rst b/docs/source/simulation_structure.rst index 2f0a56e8..6e0ab5ce 100644 --- a/docs/source/simulation_structure.rst +++ b/docs/source/simulation_structure.rst @@ -12,7 +12,7 @@ and a domain controller for managing software and users. Each node of the simulation 'tree' has responsibility for creating, deleting, and updating its direct descendants. Also, when a component's ``describe_state()`` method is called, it will include the state of its descendants. The -``apply_action()`` method can be used to act on a component or one of its descendatnts. The diagram below shows the +``apply_request()`` method can be used to act on a component or one of its descendatnts. The diagram below shows the relationship between components. .. image:: _static/component_relationship.png @@ -25,9 +25,9 @@ relationship between components. Actions ======= Agents can interact with the simulation by using actions. Actions are standardised with the -:py:class:`primaite.simulation.core.Action` class, which just holds a reference to two special functions. +:py:class:`primaite.simulation.core.RequestType` class, which just holds a reference to two special functions. -1. The action function itself, it must accept a `request` parameters which is a list of strings that describe what the +1. The request function itself, it must accept a `request` parameters which is a list of strings that describe what the action should do. It must also accept a `context` dict which can house additional information surrounding the action. For example, the context will typically include information about which entity intiated the action. 2. A validator function. This function should return a boolean value that decides if the request is permitted or not. diff --git a/docs/source/state_system.rst b/docs/source/state_system.rst new file mode 100644 index 00000000..b8a9624e --- /dev/null +++ b/docs/source/state_system.rst @@ -0,0 +1,31 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +Simulation State +============== + +``SimComponent``s in the simulation have a method called ``describe_state`` which returns a dictionary of the state of the component. This is used to report pertinent data that could impact agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` reports not only it's own attributes in the state but also that of its child components. I.e. a computer node will report the state of its ``FileSystem``, and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling childrens' own ``describe_state`` methods. + +The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then pass the state to the agents once per simulation step. For this reason, all ``SimComponent``s must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``. + +This code snippet demonstrates how the state information is defined within the ``SimComponent`` class: + +.. code-block:: python + + class Node(SimComponent): + operating_state: NodeOperatingState = NodeOperatingState.OFF + services: Dict[str, Service] = {} + + def describe_state(self) -> Dict: + state = super().describe_state() + state["operating_state"] = self.operating_state.value + state["services"] = {uuid: svc.describe_state() for uuid, svc in self.services.items()} + return state + + class Service(SimComponent): + health_state: ServiceHealthState = ServiceHealthState.GOOD + def describe_state(self) -> Dict: + state = super().describe_state() + state["health_state"] = self.health_state.value + return state From 318c9f8c5aa49d116398829ea135d1e10d2107a6 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 11:43:11 +0100 Subject: [PATCH 44/53] Fix formatting issues --- docs/source/request_system.rst | 4 ++-- docs/source/state_system.rst | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/source/request_system.rst b/docs/source/request_system.rst index 41d8eec4..09e46380 100644 --- a/docs/source/request_system.rst +++ b/docs/source/request_system.rst @@ -28,7 +28,7 @@ Just like other aspects of SimComponent, the request typess are not managed cent Since ``restart`` is a defined request type in the service's own RequestManager, the service performs a restart. Technical Detail -================ +---------------- This system was achieved by implementing two classes, :py:class:`primaite.simulator.core.RequestType`, and :py:class:`primaite.simulator.core.RequestManager`. @@ -38,7 +38,7 @@ This system was achieved by implementing two classes, :py:class:`primaite.simula The ``RequestType`` object stores a reference to a method that executes the request, for example a node could have a request type that stores a reference to ``self.turn_on()``. Technically, this can be any callable that accepts `request, context` as it's parameters. In practice, this is often defined using ``lambda`` functions within a component's ``self._init_request_manager()`` method. Optionally, the ``RequestType`` object can also hold a validator that will permit/deny the request depending on context. ``RequestManager`` -------------- +------------------ The ``RequestManager`` object stores a mapping between strings and request types. It is responsible for processing the request and passing it down the ownership tree. Technically, the ``RequestManager`` is itself a callable that accepts `request, context` tuple, and so it can be chained with other request managers. diff --git a/docs/source/state_system.rst b/docs/source/state_system.rst index b8a9624e..de4cd093 100644 --- a/docs/source/state_system.rst +++ b/docs/source/state_system.rst @@ -5,9 +5,9 @@ Simulation State ============== -``SimComponent``s in the simulation have a method called ``describe_state`` which returns a dictionary of the state of the component. This is used to report pertinent data that could impact agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` reports not only it's own attributes in the state but also that of its child components. I.e. a computer node will report the state of its ``FileSystem``, and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling childrens' own ``describe_state`` methods. +``SimComponent`` in the simulation have a method called ``describe_state`` which returns a dictionary of the state of the component. This is used to report pertinent data that could impact agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` reports not only it's own attributes in the state but also that of its child components. I.e. a computer node will report the state of its ``FileSystem``, and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling childrens' own ``describe_state`` methods. -The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then pass the state to the agents once per simulation step. For this reason, all ``SimComponent``s must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``. +The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then pass the state to the agents once per simulation step. For this reason, all ``SimComponent`` must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``. This code snippet demonstrates how the state information is defined within the ``SimComponent`` class: From 6ac7dcf9ac5e52ea07217b597def09a99683ae0e Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 11:48:22 +0100 Subject: [PATCH 45/53] Update README.md --- README.md | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 1c274a45..9ec8164e 100644 --- a/README.md +++ b/README.md @@ -18,17 +18,13 @@ PrimAITE presents the following features: - Provision of logging to support AI evaluation and metrics gathering; -- Uses the concept of Information Exchange Requirements (IERs) to model background pattern of life and adversarial behaviour; +- Realistic network traffic simulation, including address and sending packets via internet protocols like TCP, UDP, ICMP, and others -- An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP address, destination IP address, protocol and port); +- Routers with traffic routing and firewall capabilities -- Application of IERs to the platform / system laydown adheres to the ACL ruleset; +- Integration with ARCD GATE for agent training -- Presents an OpenAI gym or RLLib interface to the environment, allowing integration with any compliant defensive agents; - -- Full capture of discrete logs relating to agent training (full system state, agent actions taken, instantaneous and average reward for every step of every episode); - -- NetworkX provides laydown visualisation capability. +- Support for multiple agents, each having their own customisable observation space, action space, and reward function definition, and either deterministic or RL-directed behaviour ## Getting Started with PrimAITE @@ -50,6 +46,7 @@ python3 -m venv .venv attrib +h .venv /s /d # Hides the .venv directory .\.venv\Scripts\activate pip install https://github.com/Autonomous-Resilient-Cyber-Defence/PrimAITE/releases/download/v2.0.0/primaite-2.0.0-py3-none-any.whl +pip install GATE/arcd_gate-0.1.0-py3-none-any.whl primaite setup ``` @@ -78,6 +75,7 @@ cd ~/primaite python3 -m venv .venv source .venv/bin/activate pip install https://github.com/Autonomous-Resilient-Cyber-Defence/PrimAITE/releases/download/v2.0.0/primaite-2.0.0-py3-none-any.whl +pip install arcd_gate-0.1.0-py3-none-any.whl primaite setup ``` @@ -122,6 +120,7 @@ source venv/bin/activate ```bash python3 -m pip install -e .[dev] +pip install arcd_gate-0.1.0-py3-none-any.whl ``` #### 6. Perform the PrimAITE setup: From a06629ed0bc8648e34934a1f2028cced69494a46 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 11:48:54 +0100 Subject: [PATCH 46/53] Bump version --- src/primaite/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/VERSION b/src/primaite/VERSION index 227cea21..a6f4248b 100644 --- a/src/primaite/VERSION +++ b/src/primaite/VERSION @@ -1 +1 @@ -2.0.0 +3.0.0a1 From d4eb499729f21f524a33f76fcefdcb4524e8442b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 12:20:23 +0100 Subject: [PATCH 47/53] Start fixing up cli and setup --- sandbox.py | 1 - src/primaite/cli.py | 58 ++------------------------------- src/primaite/config/__init__.py | 2 ++ src/primaite/config/load.py | 22 +++++++++++++ 4 files changed, 26 insertions(+), 57 deletions(-) create mode 100644 src/primaite/config/__init__.py create mode 100644 src/primaite/config/load.py diff --git a/sandbox.py b/sandbox.py index ab5e701f..c5c8ae38 100644 --- a/sandbox.py +++ b/sandbox.py @@ -6,7 +6,6 @@ from primaite.game.session import PrimaiteSession _PRIMAITE_CONFIG["log_level"] = logging.DEBUG print(PRIMAITE_PATHS.app_log_dir_path) -import itertools import yaml diff --git a/src/primaite/cli.py b/src/primaite/cli.py index 9bdc414d..ea0247ec 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -10,7 +10,6 @@ import yaml from typing_extensions import Annotated from primaite import PRIMAITE_PATHS -from primaite.data_viz import PlotlyTemplate app = typer.Typer() @@ -81,14 +80,6 @@ def log_level(level: Annotated[Optional[_LogLevel], typer.Argument()] = None) -> print(f"PrimAITE Log Level: {level}") -@app.command() -def notebooks() -> None: - """Start Jupyter Lab in the users PrimAITE notebooks directory.""" - from primaite.notebooks import start_jupyter_session - - start_jupyter_session() - - @app.command() def version() -> None: """Get the installed PrimAITE version number.""" @@ -97,14 +88,6 @@ def version() -> None: print(primaite.__version__) -@app.command() -def clean_up() -> None: - """Cleans up left over files from previous version installations.""" - from primaite.setup import old_installation_clean_up - - old_installation_clean_up.run() - - @app.command() def setup(overwrite_existing: bool = True) -> None: """ @@ -113,7 +96,7 @@ def setup(overwrite_existing: bool = True) -> None: WARNING: All user-data will be lost. """ from primaite import getLogger - from primaite.setup import old_installation_clean_up, reset_demo_notebooks, reset_example_configs + from primaite.setup import reset_demo_notebooks, reset_example_configs _LOGGER = getLogger(__name__) @@ -130,19 +113,12 @@ def setup(overwrite_existing: bool = True) -> None: _LOGGER.info("Rebuilding the example notebooks...") reset_example_configs.run(overwrite_existing=True) - _LOGGER.info("Performing a clean-up of previous PrimAITE installations...") - old_installation_clean_up.run() - _LOGGER.info("PrimAITE setup complete!") @app.command() def session( - tc: Optional[str] = None, - ldc: Optional[str] = None, - load: Optional[str] = None, - legacy_tc: bool = False, - legacy_ldc: bool = False, + config: Optional[str] = None, ) -> None: """ Run a PrimAITE session. @@ -165,13 +141,8 @@ def session( legacy_ldf: If the lay down config file is a legacy file from PrimAITE < 2.0. """ from primaite.config.lay_down_config import dos_very_basic_config_path - from primaite.config.training_config import main_training_config_path from primaite.main import run - if load is not None: - # run a loaded session - run(session_path=load) - else: # start a new session using tc and ldc if not tc: @@ -186,28 +157,3 @@ def session( legacy_training_config=legacy_tc, legacy_lay_down_config=legacy_ldc, ) - - -@app.command() -def plotly_template(template: Annotated[Optional[PlotlyTemplate], typer.Argument()] = None) -> None: - """ - View or set the plotly template for Session plots. - - To View, simply call: primaite plotly-template - - To set, call: primaite plotly-template - - For example, to set as plotly_dark, call: primaite plotly-template PLOTLY_DARK - """ - if PRIMAITE_PATHS.app_config_file_path.exists(): - with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: - primaite_config = yaml.safe_load(file) - - if template: - primaite_config["session"]["outputs"]["plots"]["template"] = template.value - with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: - yaml.dump(primaite_config, file) - print(f"PrimAITE plotly template: {template.value}") - else: - template = primaite_config["session"]["outputs"]["plots"]["template"] - print(f"PrimAITE plotly template: {template}") diff --git a/src/primaite/config/__init__.py b/src/primaite/config/__init__.py new file mode 100644 index 00000000..92f5a7d2 --- /dev/null +++ b/src/primaite/config/__init__.py @@ -0,0 +1,2 @@ +# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK +"""Configuration parameters for running experiments.""" diff --git a/src/primaite/config/load.py b/src/primaite/config/load.py new file mode 100644 index 00000000..77b76299 --- /dev/null +++ b/src/primaite/config/load.py @@ -0,0 +1,22 @@ +from pathlib import Path +from typing import Union + +import yaml + +from primaite import getLogger + +_LOGGER = getLogger(__name__) + + +def load(file_path: Union[str, Path]) -> Dict: + """ + Read a YAML file and return the contents as a dictionary. + + :param file_path: Path to the YAML file. + :type file_path: Union[str, Path] + :return: Config dictionary + :rtype: Dict + """ + + if not isinstance(file_path, Path): + file_path = Path(file_path) From 1ec1b319e5208dbde18a53cc4430051d5bbd54fd Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 12:24:22 +0100 Subject: [PATCH 48/53] Fix formatting in docs --- docs/source/request_system.rst | 2 +- .../simulation_components/system/database_client_server.rst | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/request_system.rst b/docs/source/request_system.rst index 09e46380..cdaf2d99 100644 --- a/docs/source/request_system.rst +++ b/docs/source/request_system.rst @@ -5,7 +5,7 @@ Request System ============== -``SimComponent``s in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``. +``SimComponent`` in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``. Just like other aspects of SimComponent, the request typess are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved in the following way: diff --git a/docs/source/simulation_components/system/database_client_server.rst b/docs/source/simulation_components/system/database_client_server.rst index 32568477..53687f60 100644 --- a/docs/source/simulation_components/system/database_client_server.rst +++ b/docs/source/simulation_components/system/database_client_server.rst @@ -14,7 +14,7 @@ The ``DatabaseService`` provides a SQL database server simulation by extending t Key capabilities ^^^^^^^^^^^^^^^^ -- Initialises a SQLite database file in the ``Node``'s ``FileSystem`` upon creation. +- Initialises a SQLite database file in the ``Node`` 's ``FileSystem`` upon creation. - Handles connecting clients by maintaining a dictionary of connections mapped to session IDs. - Authenticates connections using a configurable password. - Executes SQL queries against the SQLite database. From b81c1739f8e632a0a3e4652e6d8c41aae8145995 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 14:26:52 +0100 Subject: [PATCH 49/53] Fix CLI and Session to work with new classes --- src/primaite/cli.py | 42 ++++++------------- .../config/_package_data/example_config.yaml | 0 src/primaite/config/load.py | 29 +++++++++++-- src/primaite/exceptions.py | 6 --- src/primaite/main.py | 35 +++++----------- src/primaite/utils/start_gate_server.py | 11 ++++- 6 files changed, 58 insertions(+), 65 deletions(-) rename example_config.yaml => src/primaite/config/_package_data/example_config.yaml (100%) diff --git a/src/primaite/cli.py b/src/primaite/cli.py index ea0247ec..43b97022 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -123,37 +123,19 @@ def session( """ Run a PrimAITE session. - tc: The training config filepath. Optional. If no value is passed then - example default training config is used from: - ~/primaite/2.0.0/config/example_config/training/training_config_main.yaml. - - ldc: The lay down config file path. Optional. If no value is passed then - example default lay down config is used from: - ~/primaite/2.0.0/config/example_config/lay_down/lay_down_config_3_doc_very_basic.yaml. - - load: The directory of a previous session. Optional. If no value is passed, then the session - will use the default training config and laydown config. Inversely, if a training config and laydown config - is passed while a session directory is passed, PrimAITE will load the session and ignore the training config - and laydown config. - - legacy_tc: If the training config file is a legacy file from PrimAITE < 2.0. - - legacy_ldf: If the lay down config file is a legacy file from PrimAITE < 2.0. + :param config: The path to the config file. Optional, if None, the example config will be used. + :type config: Optional[str] """ - from primaite.config.lay_down_config import dos_very_basic_config_path + from threading import Thread + + from primaite.config.load import example_config_path from primaite.main import run + from primaite.utils.start_gate_server import start_gate_server - else: - # start a new session using tc and ldc - if not tc: - tc = main_training_config_path() + server_thread = Thread(target=start_gate_server) + server_thread.start() - if not ldc: - ldc = dos_very_basic_config_path() - - run( - training_config_path=tc, - lay_down_config_path=ldc, - legacy_training_config=legacy_tc, - legacy_lay_down_config=legacy_ldc, - ) + if not config: + config = example_config_path() + print(config) + run(config_path=config) diff --git a/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml similarity index 100% rename from example_config.yaml rename to src/primaite/config/_package_data/example_config.yaml diff --git a/src/primaite/config/load.py b/src/primaite/config/load.py index 77b76299..b01eb129 100644 --- a/src/primaite/config/load.py +++ b/src/primaite/config/load.py @@ -1,12 +1,14 @@ from pathlib import Path -from typing import Union +from typing import Dict, Final, Union import yaml -from primaite import getLogger +from primaite import getLogger, PRIMAITE_PATHS _LOGGER = getLogger(__name__) +_EXAMPLE_CFG: Final[Path] = PRIMAITE_PATHS.user_config_path / "example_config" + def load(file_path: Union[str, Path]) -> Dict: """ @@ -17,6 +19,27 @@ def load(file_path: Union[str, Path]) -> Dict: :return: Config dictionary :rtype: Dict """ - if not isinstance(file_path, Path): file_path = Path(file_path) + if not file_path.exists(): + _LOGGER.error(f"File does not exist: {file_path}") + raise FileNotFoundError(f"File does not exist: {file_path}") + with open(file_path, "r") as file: + config = yaml.safe_load(file) + _LOGGER.debug(f"Loaded config from {file_path}") + return config + + +def example_config_path() -> Path: + """ + Get the path to the example config. + + :return: Path to the example config. + :rtype: Path + """ + path = _EXAMPLE_CFG / "example_config.yaml" + if not path.exists(): + msg = f"Example config does not exist: {path}. Have you run `primaite setup`?" + _LOGGER.error(msg) + raise FileNotFoundError(msg) + return path diff --git a/src/primaite/exceptions.py b/src/primaite/exceptions.py index 025f6d41..6aa140ba 100644 --- a/src/primaite/exceptions.py +++ b/src/primaite/exceptions.py @@ -5,12 +5,6 @@ class PrimaiteError(Exception): pass -class RLlibAgentError(PrimaiteError): - """Raised when there is a generic error with a RLlib agent that is specific to PRimAITE.""" - - pass - - class NetworkError(PrimaiteError): """Raised when an error occurs at the network level.""" diff --git a/src/primaite/main.py b/src/primaite/main.py index 0cbcff0e..831419d4 100644 --- a/src/primaite/main.py +++ b/src/primaite/main.py @@ -5,6 +5,8 @@ from pathlib import Path from typing import Optional, Union from primaite import getLogger +from primaite.config.load import load +from primaite.game.session import PrimaiteSession # from primaite.primaite_session import PrimaiteSession @@ -12,11 +14,7 @@ _LOGGER = getLogger(__name__) def run( - training_config_path: Optional[Union[str, Path]] = "", - lay_down_config_path: Optional[Union[str, Path]] = "", - session_path: Optional[Union[str, Path]] = None, - legacy_training_config: bool = False, - legacy_lay_down_config: bool = False, + config_path: Optional[Union[str, Path]] = "", ) -> None: """ Run the PrimAITE Session. @@ -32,28 +30,17 @@ def run( :param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0, otherwise False. """ - # session = PrimaiteSession( - # training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config - # ) - - # session.setup() - # session.learn() - # session.evaluate() - return NotImplemented + cfg = load(config_path) + sess = PrimaiteSession.from_config(cfg=cfg) + sess.start_session() if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--tc") - parser.add_argument("--ldc") - parser.add_argument("--load") + parser.add_argument("--config") args = parser.parse_args() - if args.load: - run(session_path=args.load) - else: - if not args.tc: - _LOGGER.error("Please provide a training config file using the --tc " "argument") - if not args.ldc: - _LOGGER.error("Please provide a lay down config file using the --ldc " "argument") - run(training_config_path=args.tc, lay_down_config_path=args.ldc) + if not args.config: + _LOGGER.error("Please provide a config file using the --config " "argument") + + run(session_path=args.config) diff --git a/src/primaite/utils/start_gate_server.py b/src/primaite/utils/start_gate_server.py index 53508cd2..d91952f2 100644 --- a/src/primaite/utils/start_gate_server.py +++ b/src/primaite/utils/start_gate_server.py @@ -1,5 +1,12 @@ """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() + +def start_gate_server(): + """Start the gate server.""" + service = GATEService() + service.start() + + +if __name__ == "__main__": + start_gate_server() From 75d1fd20c3967a8bf481d2fce319ca39fd74192b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 14:41:53 +0100 Subject: [PATCH 50/53] Remove temporary files --- sandbox.py | 21 ------ src/primaite/notebooks/.gitkeep | 0 src/primaite/notebooks/scratch.ipynb | 107 --------------------------- 3 files changed, 128 deletions(-) delete mode 100644 sandbox.py create mode 100644 src/primaite/notebooks/.gitkeep delete mode 100644 src/primaite/notebooks/scratch.ipynb diff --git a/sandbox.py b/sandbox.py deleted file mode 100644 index c5c8ae38..00000000 --- a/sandbox.py +++ /dev/null @@ -1,21 +0,0 @@ -# flake8: noqa -import logging - -from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS -from primaite.game.session import PrimaiteSession - -_PRIMAITE_CONFIG["log_level"] = logging.DEBUG -print(PRIMAITE_PATHS.app_log_dir_path) - -import yaml - -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() diff --git a/src/primaite/notebooks/.gitkeep b/src/primaite/notebooks/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/notebooks/scratch.ipynb b/src/primaite/notebooks/scratch.ipynb deleted file mode 100644 index 4e873460..00000000 --- a/src/primaite/notebooks/scratch.ipynb +++ /dev/null @@ -1,107 +0,0 @@ -{ - "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": "markdown", - "metadata": {}, - "source": [ - "### set up some services to test if actions are working" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "db_serv = net.get_node_by_hostname('database_server')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from primaite.simulator.system.services.database_service import DatabaseService" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "db_svc = DatabaseService(file_system=db_serv.file_system)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "db_serv.install_service(db_svc)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "db_serv.describe_state()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "\n" - ] - } - ], - "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 a37ce051c5659a057176b0c7d3139dedcc0a42a9 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 27 Oct 2023 14:50:31 +0100 Subject: [PATCH 51/53] Add ARCD GATE setup to primaite setup --- src/primaite/cli.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/primaite/cli.py b/src/primaite/cli.py index 43b97022..a5b3be46 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -95,6 +95,8 @@ def setup(overwrite_existing: bool = True) -> None: WARNING: All user-data will be lost. """ + from arcd_gate.cli import setup as gate_setup + from primaite import getLogger from primaite.setup import reset_demo_notebooks, reset_example_configs @@ -113,6 +115,9 @@ def setup(overwrite_existing: bool = True) -> None: _LOGGER.info("Rebuilding the example notebooks...") reset_example_configs.run(overwrite_existing=True) + _LOGGER.info("Setting up ARCD GATE...") + gate_setup() + _LOGGER.info("PrimAITE setup complete!") From e0d694a7362767f13279db28a9f1759f6cace502 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Tue, 7 Nov 2023 14:00:08 +0000 Subject: [PATCH 52/53] #2025: added GATE installation as part of pipeline --- .azure/azure-ci-build-pipeline.yaml | 11 +++++++++++ src/primaite/game/agent/observations.py | 2 +- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 9070270a..2b1af5f4 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -81,6 +81,17 @@ stages: displayName: 'Install PrimAITE' condition: eq( variables['Agent.OS'], 'Windows_NT' ) + - script: | + GATE_WHEEL=$(ls ./GATE/arcd_gate*.whl) + python -m pip install GATE_WHEEL[dev] + displayName: 'Install GATE' + condition: or(eq( variables['Agent.OS'], 'Linux' ), eq( variables['Agent.OS'], 'Darwin' )) + + - script: | + forfiles /p GATE\ /m *.whl /c "cmd /c python -m pip install @file[dev]" + displayName: 'Install GATE' + condition: eq( variables['Agent.OS'], 'Windows_NT' ) + - script: | primaite setup displayName: 'Perform PrimAITE Setup' diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 8eb322bd..a3bafeea 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -585,7 +585,7 @@ class AclObservation(AbstractObservation): self, node_ip_to_id: Dict[str, int], ports: List[int], - protocols: list[str], + protocols: List[str], where: Optional[Tuple[str]] = None, num_rules: int = 10, ) -> None: From d4679bb0e3f42bdcca5b5172ed1a91e787c3f262 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Tue, 7 Nov 2023 14:09:22 +0000 Subject: [PATCH 53/53] #2025: Missing character for GATE path --- .azure/azure-ci-build-pipeline.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 2b1af5f4..efeba284 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -83,7 +83,7 @@ stages: - script: | GATE_WHEEL=$(ls ./GATE/arcd_gate*.whl) - python -m pip install GATE_WHEEL[dev] + python -m pip install $GATE_WHEEL[dev] displayName: 'Install GATE' condition: or(eq( variables['Agent.OS'], 'Linux' ), eq( variables['Agent.OS'], 'Darwin' ))