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

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting information out of PrimAITE\n",
    "\n",
    "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!primaite setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import yaml\n",
    "from primaite import PRIMAITE_CONFIG\n",
    "from primaite.config.load import data_manipulation_config_path\n",
    "from primaite.session.environment import PrimaiteGymEnv\n",
    "from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
    "from notebook.services.config import ConfigManager\n",
    "\n",
    "cm = ConfigManager().update('notebook', {'limit_output': 50})  # limit output lines to 50 - for neatness\n",
    "\n",
    "# create the env\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "\n",
    "env = PrimaiteGymEnv(env_config=cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook opts to use the [data manipulation scenario](./Data-Manipulation-E2E-Demonstration.ipynb) (also known as UC2) network configuration but all of the methods demonstrated are config agnostic and can be used in any PrimAITE scenario."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualising the Simulation Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some of the simpler networks can be visualised by using the `.draw()` method as show in the code snippet below. \n",
    "\n",
    "Larger networks will still render but users may prefer to create their own network diagrams via other tooling as shown in the [UC7 notebooks](./UC7-E2E-Demo.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.game.simulation.network.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the state of a simulation object\n",
    "\n",
    "The state of the simulation object is used to determine the observation space used by agents.\n",
    "\n",
    "Any object created using the ``SimComponent`` class has a ``describe_state`` method which can show the state of the object.\n",
    "\n",
    "An example of such an object is ``Computer`` which inherits from ``SimComponent``. As this notebook utilises the [UC2 network configuration]((./Data-Manipulation-E2E-Demonstration.ipynb)) we can initialise the **client_1** node and confirm that it's a ``Computer`` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
    "client_1.describe_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More specific describe_state\n",
    "\n",
    "As you can see, the output from the ``describe_state`` method for the ``Computer`` object includes the describe state for all its component which causes a rather large output.\n",
    "\n",
    "As stated, the ``describe_state`` can be called on any object that inherits ``SimComponent``. Meaning, we can narrow down our output by retrieving the state of a specific item. The code snippet below calls the `describe_state` method on **client_1**'s filesystem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_1.file_system.describe_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's worth reiterating that the `describe_state()` method can be used on literally any object that inherits from ``SimComponent``. For example, even the system default **'admin'** user inherits ``.describe_state()``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_1.user_manager.admins.get(\"admin\").describe_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## System Logs \n",
    "\n",
    "Objects that inherit from the ``Node`` class will inherit the ``sys_log`` attribute.\n",
    "\n",
    "This is to simulate the idea that items such as Computer, Routers, Servers, etc. have a logging system used to diagnose problems.\n",
    "\n",
    "Enabling this functionality will slow down training time due to the amount of sheer amount of logs created therefore it's recommended to disable these logs when training/evaluating agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# store config\n",
    "# this is to prevent the notebook from breaking your local settings\n",
    "was_enabled = PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"]\n",
    "was_syslogs_enabled = PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"]\n",
    "\n",
    "# enable dev mode so that the default config outputs are overridden for this demo\n",
    "PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = True\n",
    "PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = True\n",
    "\n",
    "# Remake the environment\n",
    "env = PrimaiteGymEnv(env_config=cfg)\n",
    "\n",
    "# get the example computer\n",
    "client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
    "\n",
    "# show sys logs on terminal\n",
    "client_1.sys_log.show()\n",
    "\n",
    "# restore config\n",
    "PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = was_enabled\n",
    "PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = was_syslogs_enabled"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Viewing Agent history\n",
    "\n",
    "It's possible to view the actions carried out by an agent for a given training session using the `show_history()` method. By default, this will be all actions apart from DONOTHING actions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "\n",
    "env = PrimaiteGymEnv(env_config=cfg)\n",
    "\n",
    "# Run the training session to generate some resultant data.\n",
    "for i in range(100):\n",
    "    env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `.show_history()` should show us when the Data Manipulation used the `NODE_APPLICATION_EXECUTE` action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attacker = env.game.agents[\"data_manipulation_attacker\"]\n",
    "\n",
    "attacker.show_history()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}