{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customising UC2 Red Agents\n",
    "\n",
    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
    "\n",
    "This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
    "\n",
    "First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
    "\n",
    "*(For a full explanation of the Data Manipulation scenario, check out the data manipulation scenario notebook)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!primaite setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "\n",
    "from primaite.config.load import data_manipulation_config_path\n",
    "from primaite.game.agent.interface import AgentHistoryItem\n",
    "from primaite.session.environment import PrimaiteGymEnv\n",
    "import yaml\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_cfg_have_flat_obs(cfg):\n",
    "    for agent in cfg['agents']:\n",
    "        if agent['type'] == \"ProxyAgent\":\n",
    "            agent['agent_settings']['flatten_obs'] = False"
   ]
  },
  {
   "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",
    "    make_cfg_have_flat_obs(cfg)\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "obs, info = env.reset()\n",
    "print('env created successfully')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def friendly_output_red_action(info):\n",
    "    # parse the info dict form step output and write out what the red agent is doing\n",
    "    red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
    "    red_action = red_info.action\n",
    "    if red_action == 'do_nothing':\n",
    "        red_str = 'DO NOTHING'\n",
    "    elif red_action == 'node_application_execute':\n",
    "        client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
    "        red_str = f\"ATTACK from {client}\"\n",
    "    return red_str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the red agent can start on client 1 or client 2. It starts its attack on a random step between 20 and 30, and it repeats its attack every 15-25 steps.\n",
    "\n",
    "It also has a 20% chance to fail to perform the port scan, and a 20% chance to fail launching the SQL attack. However it will continue where it left off after a failed step. I.e. if lucky, it can perform the port scan and SQL attack on the first try. If the port scan works, but the sql attack fails the first time it tries to attack, the next time it will not need to port scan again, it can go straight to trying to use SQL attack again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(35):\n",
    "    step_num = env.game.step_counter\n",
    "    obs, reward, terminated, truncated, info = env.step(0)\n",
    "    red = friendly_output_red_action(info)\n",
    "    print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the agent does nothing most of the time, let's only print the steps where it performs an attack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()\n",
    "for step in range(100):\n",
    "    step_num = env.game.step_counter\n",
    "    obs, reward, terminated, truncated, info = env.step(0)\n",
    "    red = friendly_output_red_action(info)\n",
    "    if red.startswith(\"ATTACK\"):\n",
    "        print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Red Configuration\n",
    "\n",
    "There are two important parts of the YAML config for varying red agent behaviour.\n",
    "\n",
    "### Red agent settings\n",
    "Here is an annotated config for the red agent in the data manipulation scenario."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```yaml\n",
    "  - ref: data_manipulation_attacker  # name of agent\n",
    "    team: RED # not used, just for human reference\n",
    "    type: RedDatabaseCorruptingAgent  # type of agent - this lets primaite know which agent class to use\n",
    "\n",
    "    # Since the agent does not need to react to what is happening in the environment, the observation space is empty.\n",
    "    observation_space:\n",
    "      type: UC2RedObservation\n",
    "      options:\n",
    "        nodes: {}\n",
    "\n",
    "    action_space:\n",
    "    \n",
    "      # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
    "      options:\n",
    "        nodes:\n",
    "        - node_name: client_1 # The network should have a node called client_1\n",
    "          applications:\n",
    "            - application_name: DataManipulationBot # The node client_1 should have DataManipulationBot configured on it\n",
    "        - node_name: client_2 # The network should have a node called client_2\n",
    "          applications:\n",
    "            - application_name: DataManipulationBot # The node client_2 should have DataManipulationBot configured on it\n",
    "\n",
    "        # not important\n",
    "        max_folders_per_node: 1\n",
    "        max_files_per_folder: 1\n",
    "        max_services_per_node: 1\n",
    "\n",
    "    # red agent does not need a reward function\n",
    "    reward_function:\n",
    "      reward_components:\n",
    "        - type: DUMMY\n",
    "\n",
    "    # These actions are passed to the RedDatabaseCorruptingAgent init method, they dictate the schedule of attacks\n",
    "    agent_settings:\n",
    "      start_settings:\n",
    "        start_step: 25  # first attack at step 25\n",
    "        frequency: 20  # attacks will happen every 20 steps (on average)\n",
    "        variance: 5  # the timing of attacks will vary by up to 5 steps earlier or later\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Malicious application settings\n",
    "The red agent uses an application called `DataManipulationBot` which leverages a node's `DatabaseClient` to send a malicious SQL query to the database server. Here's an annotated example of how this is configured in the yaml *(with impertinent config items omitted)*:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```yaml\n",
    "simulation:\n",
    "  network:\n",
    "    nodes:\n",
    "    - ref: client_1\n",
    "      hostname: client_1\n",
    "      type: computer\n",
    "      ip_address: 192.168.10.21\n",
    "      subnet_mask: 255.255.255.0\n",
    "      default_gateway: 192.168.10.1\n",
    "      \n",
    "      # \n",
    "      applications:\n",
    "      - ref: data_manipulation_bot\n",
    "        type: DataManipulationBot\n",
    "        options:\n",
    "          port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
    "          data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
    "          payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
    "          server_ip: 192.168.1.14 # IP address of server hosting the database\n",
    "      - ref: client_1_database_client\n",
    "        type: DatabaseClient # Database client must be installed in order for DataManipulationBot to function\n",
    "        options:\n",
    "          db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Editing red agent settings\n",
    "\n",
    "### Removing randomness from attack timing\n",
    "\n",
    "We can make the attacks happen at completely predictable intervals if we edit the red agent's settings to set variance to 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "change = yaml.safe_load(\"\"\"\n",
    "start_settings:\n",
    "  start_step: 25\n",
    "  frequency: 20\n",
    "  variance: 0\n",
    "\"\"\")\n",
    "\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    for agent in cfg['agents']:\n",
    "      if agent['ref'] == \"data_manipulation_attacker\":\n",
    "        agent['agent_settings'] = change\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "env.reset()\n",
    "for step in range(100):\n",
    "    step_num = env.game.step_counter\n",
    "    obs, reward, terminated, truncated, info = env.step(0)\n",
    "    red = friendly_output_red_action(info)\n",
    "    if red.startswith(\"ATTACK\"):\n",
    "        print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making the start node always the same\n",
    "\n",
    "Normally, the agent randomly chooses between the nodes in its action space to send attacks from:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Open the config without changing anything\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "env.reset()\n",
    "for ep in range(12):\n",
    "    env.reset()\n",
    "    for step in range(31):\n",
    "        step_num = env.game.step_counter\n",
    "        obs, reward, terminated, truncated, info = env.step(0)\n",
    "        red = friendly_output_red_action(info)\n",
    "        if red.startswith(\"ATTACK\"):\n",
    "            print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make the agent always start on a node of our choice letting that be the only node in the agent's action space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "change = yaml.safe_load(\"\"\"\n",
    "# TODO:\n",
    "\"\"\")\n",
    "#TODO 2869 fix\n",
    "\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    for agent in cfg['agents']:\n",
    "      if agent['ref'] == \"data_manipulation_attacker\":\n",
    "        agent.update(change)\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "env.reset()\n",
    "for ep in range(12):\n",
    "    env.reset()\n",
    "    for step in range(31):\n",
    "        step_num = env.game.step_counter\n",
    "        obs, reward, terminated, truncated, info = env.step(0)\n",
    "        red = friendly_output_red_action(info)\n",
    "        if red.startswith(\"ATTACK\"):\n",
    "            print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make the attack less likely to succeed.\n",
    "\n",
    "We can change the success probabilities within the data manipulation bot application. When the attack succeeds, the reward goes down.\n",
    "\n",
    "Setting the probabilities to 1.0 means the attack always succeeds - the reward will always drop\n",
    "\n",
    "Setting the probabilities to 0.0 means the attack always fails - the reward will never drop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make attack always succeed.\n",
    "change = yaml.safe_load(\"\"\"\n",
    "      applications:\n",
    "      - ref: data_manipulation_bot\n",
    "        type: DataManipulationBot\n",
    "        options:\n",
    "          port_scan_p_of_success: 1.0\n",
    "          data_manipulation_p_of_success: 1.0\n",
    "          payload: \"DELETE\"\n",
    "          server_ip: 192.168.1.14\n",
    "      - ref: client_1_web_browser\n",
    "        type: WebBrowser\n",
    "        options:\n",
    "          target_url: http://arcd.com/users/\n",
    "      - ref: client_1_database_client\n",
    "        type: DatabaseClient\n",
    "        options:\n",
    "          db_server_ip: 192.168.1.14\n",
    "\"\"\")\n",
    "\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    cfg['simulation']['network']\n",
    "    for node in cfg['simulation']['network']['nodes']:\n",
    "      if node['hostname'] in ['client_1', 'client_2']:\n",
    "        node['applications'] = change['applications']\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "env.reset()\n",
    "for ep in range(5):\n",
    "    env.reset()\n",
    "    for step in range(36):\n",
    "        step_num = env.game.step_counter\n",
    "        obs, reward, terminated, truncated, info = env.step(0)\n",
    "        red = friendly_output_red_action(info)\n",
    "        if step_num == 35:\n",
    "            print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make attack always fail.\n",
    "change = yaml.safe_load(\"\"\"\n",
    "      applications:\n",
    "      - ref: data_manipulation_bot\n",
    "        type: DataManipulationBot\n",
    "        options:\n",
    "          port_scan_p_of_success: 0.0\n",
    "          data_manipulation_p_of_success: 0.0\n",
    "          payload: \"DELETE\"\n",
    "          server_ip: 192.168.1.14\n",
    "      - ref: client_1_web_browser\n",
    "        type: WebBrowser\n",
    "        options:\n",
    "          target_url: http://arcd.com/users/\n",
    "      - ref: client_1_database_client\n",
    "        type: DatabaseClient\n",
    "        options:\n",
    "          db_server_ip: 192.168.1.14\n",
    "\"\"\")\n",
    "\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    cfg['simulation']['network']\n",
    "    for node in cfg['simulation']['network']['nodes']:\n",
    "      if node['hostname'] in ['client_1', 'client_2']:\n",
    "        node['applications'] = change['applications']\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "env.reset()\n",
    "for ep in range(5):\n",
    "    env.reset()\n",
    "    for step in range(36):\n",
    "        step_num = env.game.step_counter\n",
    "        obs, reward, terminated, truncated, info = env.step(0)\n",
    "        red = friendly_output_red_action(info)\n",
    "        if step_num == 35:\n",
    "            print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
