{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Command and Control Application Suite E2E Demonstration\n",
    "\n",
    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
    "\n",
    "This notebook demonstrates the current implementation of the command and control (C2) server and beacon applications in primAITE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import yaml\n",
    "from primaite.config.load import data_manipulation_config_path\n",
    "from primaite.session.environment import PrimaiteGymEnv\n",
    "from primaite.simulator.network.hardware.nodes.network.router import Router\n",
    "from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon\n",
    "from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server\n",
    "from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import C2Command\n",
    "from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n",
    "from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
    "from primaite.simulator.network.hardware.nodes.host.server import Server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Notebook Setup** | **Network Configuration:**\n",
    "\n",
    "This notebook uses the same network setup as UC2. Please refer to the main [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb).\n",
    "\n",
    "However, this notebook replaces the red agent used in UC2 with a custom proxy red agent built for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_c2_agent = \"\"\"\n",
    "  - ref: CustomC2Agent\n",
    "    team: RED\n",
    "    type: ProxyAgent\n",
    "\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes:\n",
    "          - node_name: web_server\n",
    "            applications:\n",
    "              - application_name: C2Beacon\n",
    "          - node_name: client_1\n",
    "            applications:\n",
    "              - application_name: C2Server\n",
    "        max_folders_per_node: 1\n",
    "        max_files_per_folder: 1\n",
    "        max_services_per_node: 2\n",
    "        max_nics_per_node: 8\n",
    "        max_acl_rules: 10\n",
    "        ip_list:\n",
    "          - 192.168.1.21\n",
    "          - 192.168.1.14\n",
    "        wildcard_list:\n",
    "          - 0.0.0.1\n",
    "      action_map:\n",
    "        0:\n",
    "          action: do_nothing\n",
    "          options: {}\n",
    "        1:\n",
    "          action: node_application_install\n",
    "          options:\n",
    "            node_id: 0\n",
    "            application_name: C2Beacon\n",
    "        2:\n",
    "          action: configure_c2_beacon\n",
    "          options:\n",
    "            node_id: 0\n",
    "            config:\n",
    "              c2_server_ip_address: 192.168.10.21\n",
    "              keep_alive_frequency:\n",
    "              masquerade_protocol:\n",
    "              masquerade_port:\n",
    "        3:\n",
    "          action: node_application_execute\n",
    "          options:\n",
    "            node_id: 0\n",
    "            application_id: 0\n",
    "        4:\n",
    "          action: c2_server_terminal_command\n",
    "          options:\n",
    "            node_id: 1\n",
    "            ip_address:\n",
    "            account:\n",
    "              username: admin\n",
    "              password: admin\n",
    "            commands:\n",
    "              -\n",
    "                - software_manager\n",
    "                - application\n",
    "                - install\n",
    "                - RansomwareScript\n",
    "        5:\n",
    "          action: c2_server_ransomware_configure\n",
    "          options:\n",
    "            node_id: 1\n",
    "            config:\n",
    "              server_ip_address: 192.168.1.14\n",
    "              payload: ENCRYPT\n",
    "        6:\n",
    "          action: c2_server_data_exfiltrate\n",
    "          options:\n",
    "            node_id: 1\n",
    "            target_file_name: \"database.db\"\n",
    "            target_folder_name: \"database\"\n",
    "            exfiltration_folder_name: \"spoils\"\n",
    "            target_ip_address: 192.168.1.14\n",
    "            account:\n",
    "              username: admin\n",
    "              password: admin\n",
    "\n",
    "        7:\n",
    "          action: c2_server_ransomware_launch\n",
    "          options:\n",
    "            node_id: 1\n",
    "        8:\n",
    "          action: configure_c2_beacon\n",
    "          options:\n",
    "            node_id: 0\n",
    "            config:\n",
    "              c2_server_ip_address: 192.168.10.21\n",
    "              keep_alive_frequency: 10\n",
    "              masquerade_protocol: TCP\n",
    "              masquerade_port: DNS\n",
    "        9:\n",
    "          action: configure_c2_beacon\n",
    "          options:\n",
    "            node_id: 0\n",
    "            config:\n",
    "              c2_server_ip_address: 192.168.10.22\n",
    "              keep_alive_frequency:\n",
    "              masquerade_protocol:\n",
    "              masquerade_port:\n",
    "\n",
    "    reward_function:\n",
    "      reward_components:\n",
    "        - type: DUMMY\n",
    "\"\"\"\n",
    "c2_agent_yaml = yaml.safe_load(custom_c2_agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_manipulation_config_path()) as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    # removing all agents & adding the custom agent.\n",
    "    cfg['agents'] = {}\n",
    "    cfg['agents'] = c2_agent_yaml\n",
    "\n",
    "\n",
    "env = PrimaiteGymEnv(env_config=cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Notebook Setup** | Network Prerequisites\n",
    "\n",
    "Before the Red Agent is able to perform any C2 specific actions, the C2 Server needs to be installed and run.\n",
    "This is because in higher fidelity environments (and the real-world) a C2 server would not be accessible by a private network blue agent and the C2 Server would already be in place before the an adversary (Red Agent) starts.\n",
    "\n",
    "The cells below install and run the C2 Server on client_1 directly via the simulation API."
   ]
  },
  {
   "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.software_manager.install(C2Server)\n",
    "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "c2_server.run()\n",
    "client_1.software_manager.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Command and Control** | C2 Beacon Actions\n",
    "\n",
    "Before a C2 Server can accept any commands it must first establish connection with a C2 Beacon.\n",
    "\n",
    "A red agent is able to install, configure and establish a C2 beacon at any point in an episode. The code cells below demonstrate the actions and option parameters that are needed to perform this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Beacon Actions | node_application_install\n",
    "\n",
    "The custom proxy red agent defined at the start of this notebook has been configured to install the C2 Beacon as action ``1`` in it's action map. \n",
    "\n",
    "The below yaml snippet shows all the relevant agent options for this action:\n",
    "\n",
    "```yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "          - node_name: web_server\n",
    "            applications: \n",
    "              - application_name: C2Beacon\n",
    "        ...\n",
    "    ...\n",
    "      action_map:\n",
    "        1:\n",
    "          action: node_application_install \n",
    "          options:\n",
    "            node_id: 0 # Index 0 at the node list.\n",
    "            application_name: C2Beacon\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(1)\n",
    "web_server: Computer = env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "web_server.software_manager.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Beacon Actions | configure_c2_beacon \n",
    "\n",
    "The custom proxy red agent defined at the start of this notebook can configure the C2 Beacon via action ``2`` in it's action map. \n",
    "\n",
    "The yaml snippet below shows all the relevant agent options for this action:\n",
    "\n",
    "```yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "          - node_name: web_server\n",
    "        ...\n",
    "    ...\n",
    "      action_map:\n",
    "        ...\n",
    "        2:\n",
    "          action: configure_c2_beacon\n",
    "          options:\n",
    "            node_id: 0 # Node Index\n",
    "            config: # Further information about these config options can be found at the bottom of this notebook.\n",
    "              c2_server_ip_address: 192.168.10.21\n",
    "              keep_alive_frequency:\n",
    "              masquerade_protocol:\n",
    "              masquerade_port:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(2)\n",
    "c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
    "web_server.software_manager.show()\n",
    "c2_beacon.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Beacon Actions | node_application_execute\n",
    "\n",
    "The final action is ``node_application_execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n",
    "\n",
    "The yaml snippet below shows all the relevant agent options for this action:\n",
    "\n",
    "```yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "          - node_name: web_server\n",
    "              applications: \n",
    "              - application_name: C2Beacon\n",
    "        ...\n",
    "    ...\n",
    "      action_map:\n",
    "        ...\n",
    "        3:\n",
    "          action: node_application_execute\n",
    "          options:\n",
    "            node_id: 0\n",
    "            application_id: 0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_beacon.show()\n",
    "c2_server.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Command and Control** | C2 Server Actions\n",
    "\n",
    "Once the C2 suite has been successfully established, the C2 Server based actions become available to the Red Agent. \n",
    "\n",
    "\n",
    "This next section will demonstrate the different actions that become available to a red agent after establishing a C2 connection:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Server Actions | C2_SERVER_TERMINAL_COMMAND\n",
    "\n",
    "The C2 Server's terminal action: ``C2_SERVER_TERMINAL_COMMAND`` is indexed at  ``4`` in it's action map. \n",
    "\n",
    "This action leverages the terminal service that is installed by default on all nodes to grant red agents a lot more configurability. If you're unfamiliar with terminals then it's recommended that you refer to the ``Terminal Processing`` notebook.\n",
    "\n",
    "It's worth noting that an additional benefit a red agent has when using the terminal service via the C2 Server is that you can execute multiple commands in one action. \n",
    "\n",
    "In this notebook, the ``C2_SERVER_TERMINAL_COMMAND`` is used to install a RansomwareScript application on the ``web_server`` node.\n",
    "\n",
    "The yaml snippet below shows all the relevant agent options for this action:\n",
    "\n",
    "``` yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "        ...\n",
    "        - node_name: client_1\n",
    "            applications: \n",
    "              - application_name: C2Server\n",
    "        ...\n",
    "        action_map:\n",
    "           4:\n",
    "            action: C2_SERVER_TERMINAL_COMMAND\n",
    "            options:\n",
    "              node_id: 1\n",
    "              ip_address:\n",
    "              account:\n",
    "              username: admin\n",
    "              password: admin\n",
    "              commands:\n",
    "              - \n",
    "                  - software_manager\n",
    "                  - application\n",
    "                  - install\n",
    "                  - RansomwareScript\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_1.software_manager.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Server Actions | c2_server_ransomware_configure\n",
    "\n",
    "Another action the C2 Server grants is the ability for a Red Agent to configure the RansomwareScript via the C2 Server rather than the note directly.\n",
    "\n",
    "This action is indexed as action ``5``.\n",
    "\n",
    "The yaml snippet below shows all the relevant agent options for this action:\n",
    "\n",
    "``` yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "        ...\n",
    "        - node_name: client_1\n",
    "            applications: \n",
    "              - application_name: C2Server\n",
    "        ...\n",
    "        action_map:\n",
    "           5:\n",
    "            action: c2_server_ransomware_configure\n",
    "            options:\n",
    "              node_id: 1\n",
    "              config:\n",
    "                server_ip_address: 192.168.1.14\n",
    "                payload: ENCRYPT\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ransomware_script: RansomwareScript = web_server.software_manager.software[\"RansomwareScript\"]\n",
    "web_server.software_manager.show()\n",
    "ransomware_script.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Server Actions | c2_server_data_exfiltrate\n",
    "\n",
    "The second to last action available is the ``c2_server_data_exfiltrate`` which is indexed as action ``6`` in the action map.\n",
    "\n",
    "This action can be used to exfiltrate a target file on a remote node to the C2 Beacon and the C2 Server's host file system via the ``FTP`` services.\n",
    "\n",
    "The below yaml snippet shows all the relevant agent options for this action:\n",
    "\n",
    "``` yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "        ...\n",
    "        - node_name: client_1\n",
    "            applications: \n",
    "              - application_name: C2Server\n",
    "        ...\n",
    "        action_map:\n",
    "           6:\n",
    "            action: c2_server_data_exfiltrate\n",
    "            options:\n",
    "              node_id: 1\n",
    "              target_file_name: \"database.db\"\n",
    "              target_folder_name: \"database\"\n",
    "              exfiltration_folder_name: \"spoils\"\n",
    "              target_ip_address: \"192.168.1.14\"\n",
    "              account:\n",
    "                username: \"admin\",\n",
    "                password: \"admin\"\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(6)"
   ]
  },
  {
   "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.software_manager.file_system.show(full=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "web_server: Computer = env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "web_server.software_manager.file_system.show(full=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | C2 Server Actions | c2_server_ransomware_launch\n",
    "\n",
    "Finally, the last available action is for the c2_server_ransomware_launch to start the ransomware script installed on the same node as the C2 beacon.\n",
    "\n",
    "This action is indexed as action ``7``.\n",
    "\n",
    "\"The yaml snippet below shows all the relevant agent options for this action:\n",
    "\n",
    "``` yaml\n",
    "    action_space:\n",
    "      options:\n",
    "        nodes: # Node List\n",
    "        ...\n",
    "        - node_name: client_1\n",
    "            applications: \n",
    "              - application_name: C2Server\n",
    "        ...\n",
    "        action_map:\n",
    "           7:\n",
    "            action: c2_server_ransomware_launch\n",
    "            options:\n",
    "              node_id: 1\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "database_server: Server = env.game.simulation.network.get_node_by_hostname(\"database_server\")\n",
    "database_server.software_manager.file_system.show(full=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Command and Control** | Blue Agent Relevance\n",
    "\n",
    "The next section of the notebook will demonstrate the impact the command and control suite has on the Blue Agent's observation space as well as some potential actions that can be used to prevent the attack from being successful.\n",
    "\n",
    "The code cell below recreates the UC2 network and swaps out the previous custom red agent with a custom blue agent. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_blue_agent_yaml = \"\"\"\n",
    "  - ref: defender\n",
    "    team: BLUE\n",
    "    type: ProxyAgent\n",
    "\n",
    "    observation_space:\n",
    "      type: CUSTOM\n",
    "      options:\n",
    "        components:\n",
    "          - type: NODES\n",
    "            label: NODES\n",
    "            options:\n",
    "              hosts:\n",
    "                - hostname: web_server\n",
    "                  applications:\n",
    "                    - application_name: C2Beacon\n",
    "                    - application_name: RansomwareScript\n",
    "                  folders:\n",
    "                    - folder_name: exfiltration_folder\n",
    "                      files:\n",
    "                      - file_name: database.db\n",
    "                - hostname: database_server\n",
    "                  folders:\n",
    "                    - folder_name: exfiltration_folder\n",
    "                      files:\n",
    "                      - file_name: database.db\n",
    "                - hostname: client_1\n",
    "                - hostname: client_2\n",
    "              num_services: 0\n",
    "              num_applications: 2\n",
    "              num_folders: 1\n",
    "              num_files: 1\n",
    "              num_nics: 1\n",
    "              include_num_access: true\n",
    "              include_nmne: false\n",
    "              monitored_traffic:\n",
    "                icmp:\n",
    "                    - NONE\n",
    "                tcp:\n",
    "                    - HTTP\n",
    "                    - DNS\n",
    "                    - FTP\n",
    "              routers:\n",
    "                - hostname: router_1\n",
    "              num_ports: 3\n",
    "              ip_list:\n",
    "                - 192.168.1.10\n",
    "                - 192.168.1.12\n",
    "                - 192.168.1.14\n",
    "                - 192.168.1.16\n",
    "                - 192.168.1.110\n",
    "                - 192.168.10.21\n",
    "                - 192.168.10.22\n",
    "                - 192.168.10.110\n",
    "              wildcard_list:\n",
    "                - 0.0.0.1\n",
    "              port_list:\n",
    "                - 80\n",
    "                - 53\n",
    "                - 21\n",
    "              protocol_list:\n",
    "                - ICMP\n",
    "                - TCP\n",
    "                - UDP\n",
    "              num_rules: 10\n",
    "\n",
    "          - type: LINKS\n",
    "            label: LINKS\n",
    "            options:\n",
    "              link_references:\n",
    "                - router_1:eth-1<->switch_1:eth-8\n",
    "                - router_1:eth-2<->switch_2:eth-8\n",
    "                - switch_1:eth-1<->web_server:eth-1\n",
    "                - switch_1:eth-2<->web_server:eth-1\n",
    "                - switch_1:eth-3<->database_server:eth-1\n",
    "                - switch_1:eth-4<->backup_server:eth-1\n",
    "                - switch_1:eth-7<->security_suite:eth-1\n",
    "                - switch_2:eth-1<->client_1:eth-1\n",
    "                - switch_2:eth-2<->client_2:eth-1\n",
    "                - switch_2:eth-7<->security_suite:eth-2\n",
    "          - type: \"NONE\"\n",
    "            label: ICS\n",
    "            options: {}\n",
    "\n",
    "    action_space:\n",
    "      action_map:\n",
    "          0:\n",
    "            action: do_nothing\n",
    "            options: {}\n",
    "          1:\n",
    "            action: node_application_remove\n",
    "            options:\n",
    "              node_id: 0\n",
    "              application_name: C2Beacon\n",
    "          2:\n",
    "            action: node_shutdown\n",
    "            options:\n",
    "              node_id: 0\n",
    "          3:\n",
    "            action: router_acl_add_rule\n",
    "            options:\n",
    "              target_router: router_1\n",
    "              position: 1\n",
    "              permission: 2\n",
    "              source_ip_id: 2\n",
    "              dest_ip_id: 3\n",
    "              source_port_id: 2\n",
    "              dest_port_id: 2\n",
    "              protocol_id: 1\n",
    "              source_wildcard_id: 0\n",
    "              dest_wildcard_id: 0\n",
    "\n",
    "\n",
    "      options:\n",
    "        nodes:\n",
    "        - node_name: web_server\n",
    "          applications:\n",
    "          - application_name: C2Beacon\n",
    "\n",
    "        - node_name: database_server\n",
    "          folders:\n",
    "          - folder_name: database\n",
    "            files:\n",
    "            - file_name: database.db\n",
    "          services:\n",
    "          - service_name: DatabaseService\n",
    "        - node_name: router_1\n",
    "\n",
    "        max_folders_per_node: 2\n",
    "        max_files_per_folder: 2\n",
    "        max_services_per_node: 2\n",
    "        max_nics_per_node: 8\n",
    "        max_acl_rules: 10\n",
    "        ip_list:\n",
    "          - 192.168.10.21\n",
    "          - 192.168.1.12\n",
    "        wildcard_list:\n",
    "          - 0.0.0.1\n",
    "    reward_function:\n",
    "      reward_components:\n",
    "        - type: DUMMY\n",
    "\n",
    "    agent_settings:\n",
    "      flatten_obs: False\n",
    "\"\"\"\n",
    "custom_blue = yaml.safe_load(custom_blue_agent_yaml)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_manipulation_config_path()) as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    # removing all agents & adding the custom agent.\n",
    "    cfg['agents'] = {}\n",
    "    cfg['agents'] = custom_blue\n",
    "\n",
    "\n",
    "blue_env = PrimaiteGymEnv(env_config=cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function for showing OBS changes between each time step.\n",
    "\n",
    "from deepdiff.diff import DeepDiff\n",
    "\n",
    "def display_obs_diffs(old, new, step_counter):\n",
    "    \"\"\"\n",
    "    Use DeepDiff to extract and display differences in old and new instances of\n",
    "    the observation space.\n",
    "\n",
    "    :param old: observation space instance.\n",
    "    :param new: observation space instance.\n",
    "    :param step_counter: current step counter.\n",
    "    \"\"\"\n",
    "    print(\"\\nObservation space differences\")\n",
    "    print(\"-----------------------------\")\n",
    "    diff = DeepDiff(old, new)\n",
    "    print(f\"Step {step_counter}\")\n",
    "    for d,v in diff.get('values_changed', {}).items():\n",
    "        print(f\"{d}: {v['old_value']} -> {v['new_value']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | Blue Agent Relevance | Observation Space\n",
    "\n",
    "This section demonstrates the impacts that each of that the C2 Beacon and the C2 Server's commands cause on the observation space (OBS)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Command and Control** | OBS Impact | C2 Beacon | Installation & Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resetting the environment and capturing the default observation space.\n",
    "blue_env.reset()\n",
    "default_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the C2 Suite via the simulation API.\n",
    "\n",
    "client_1: Computer = blue_env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
    "web_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "\n",
    "# Installing the C2 Server.\n",
    "client_1.software_manager.install(C2Server)\n",
    "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "c2_server.run()\n",
    "\n",
    "# Installing the C2 Beacon.\n",
    "web_server.software_manager.install(C2Beacon)\n",
    "c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
    "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
    "c2_beacon.establish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capturing the observation impacts of the previous code cell: C2 Suite setup.\n",
    "c2_configuration_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(default_obs, c2_configuration_obs, blue_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Command and Control** | OBS Impact | C2 Server | Terminal Command\n",
    "\n",
    "Using the C2 Server's ``TERMINAL`` command it is possible to install a ``RansomwareScript`` application onto the C2 Beacon's host.\n",
    "\n",
    "The below code cells perform this as well as capturing the OBS impacts.\n",
    "\n",
    "It's important to note that the ``TERMINAL`` command is not limited to just installing software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing RansomwareScript via C2 Terminal Commands\n",
    "ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"]],\n",
    "                              \"username\": \"admin\",\n",
    "                              \"password\": \"admin\"}\n",
    "c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuring the RansomwareScript\n",
    "ransomware_config = {\"server_ip_address\": \"192.168.1.14\", \"payload\": \"ENCRYPT\"}\n",
    "c2_server.send_command(C2Command.RANSOMWARE_CONFIGURE, command_options=ransomware_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capturing the observation impacts of the previous code cell: Ransomware installation & configuration.\n",
    "c2_ransomware_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(default_obs, c2_ransomware_obs, env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Command and Control** | OBS Impact | C2 Server | Data Exfiltration\n",
    "\n",
    "Before encrypting the database.db file, the ``DATA_EXFILTRATION`` command can be used to copy the database.db file onto both the C2 Server and the C2 Beacon's file systems:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exfil_options={\n",
    "    \"username\": \"admin\",\n",
    "    \"password\": \"admin\",\n",
    "    \"target_ip_address\": \"192.168.1.14\",\n",
    "    \"target_folder_name\": \"database\",\n",
    "    \"exfiltration_folder_name\": \"exfiltration_folder\",\n",
    "    \"target_file_name\": \"database.db\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_server.send_command(given_command=C2Command.DATA_EXFILTRATION, command_options=exfil_options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_exfil_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(c2_ransomware_obs, c2_exfil_obs, env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Command and Control** | OBS Impact | C2 Server | Ransomware Commands\n",
    "\n",
    "The code cell below demonstrates the differences between the ransomware script installation obs and the impact of RansomwareScript upon the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuring the RansomwareScript\n",
    "ransomware_config = {\"server_ip_address\": \"192.168.1.14\", \"payload\": \"ENCRYPT\"}\n",
    "c2_server.send_command(C2Command.RANSOMWARE_CONFIGURE, command_options=ransomware_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Waiting for the ransomware to finish installing and then launching the RansomwareScript.\n",
    "blue_env.step(0)\n",
    "c2_server.send_command(C2Command.RANSOMWARE_LAUNCH, command_options={})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capturing the observation impacts of the previous code cell: Launching the RansomwareScript.\n",
    "c2_final_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(c2_ransomware_obs, c2_final_obs, blue_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | Blue Agent Relevance | Action Space\n",
    "\n",
    "The next section of this notebook will go over some potential blue agent actions that could be use to thwart the previously demonstrated attack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This method is used to simplify setting up the C2Server and the C2 Beacon.\n",
    "def c2_setup(given_env: PrimaiteGymEnv):\n",
    "    client_1: Computer = given_env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
    "    web_server: Server = given_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "\n",
    "    client_1.software_manager.install(C2Server)\n",
    "    c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "    c2_server.run()\n",
    "\n",
    "    web_server.software_manager.install(C2Beacon)\n",
    "    c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
    "    c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
    "    c2_beacon.establish()\n",
    "\n",
    "    return given_env, c2_server, c2_beacon, client_1, web_server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Removing the C2 Beacon.\n",
    "\n",
    "The simplest way a blue agent could prevent the C2 suite is by simply removing the C2 beacon from it's installation point. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the C2 Suite using the c2_setup method & capturing the OBS impacts\n",
    "\n",
    "blue_env, c2_server, c2_beacon, client_1, web_server = c2_setup(given_env=blue_env)\n",
    "pre_blue_action_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code cell below uses the custom blue agent defined at the start of this section perform a node_application_remove on the C2 beacon:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using CAOS ACTION: node_application_remove & capturing the OBS\n",
    "post_blue_action_obs, _, _, _, _ = blue_env.step(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can see after the effects of after stepping another timestep and looking at the web_servers software manager and the OBS differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_env.step(0)\n",
    "web_server.software_manager.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(pre_blue_action_obs, post_blue_action_obs, blue_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are unable to do so as the C2 Server has lost its connection to the C2 Beacon:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attempting to install the C2 RansomwareScript\n",
    "ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"]],\n",
    "                            \"username\": \"admin\",\n",
    "                            \"password\": \"admin\"}\n",
    "\n",
    "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Shutting down the node infected with a C2 Beacon.\n",
    "\n",
    "Another way a blue agent can prevent the C2 suite is by shutting down the C2 beacon's host node. Whilst not as effective as the previous option, depending on the situation (such as multiple malicious applications) or other scenarios it may be more timestep efficient for a blue agent to shut down a node directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the C2 Suite using the c2_setup method & capturing the OBS impacts\n",
    "\n",
    "blue_env, c2_server, c2_beacon, client_1, web_server = c2_setup(given_env=blue_env)\n",
    "pre_blue_action_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code cell below uses the custom blue agent defined at the start of this section to perform a ``node_shut_down`` action on the web server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using CAOS ACTION: node_shut_down & capturing the OBS\n",
    "post_blue_action_obs, _, _, _, _ = blue_env.step(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can see the effects of after another timestep and looking at the web server's operating state & the OBS differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "web_server = blue_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "print(web_server.operating_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(pre_blue_action_obs, post_blue_action_obs, blue_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attempting to install the C2 RansomwareScript\n",
    "ransomware_install_command = {\"commands\":[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"],\n",
    "                            \"username\": \"admin\",\n",
    "                            \"password\": \"admin\"}\n",
    "\n",
    "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Blocking C2 Traffic via ACL.\n",
    "\n",
    "Another potential option a blue agent could take is by placing an ACL rule which blocks traffic between the C2 Server and C2 Beacon.\n",
    "\n",
    "It's worth noting the potential effectiveness of this approach is connected to the current green agent traffic on the network. For example, if there are multiple green agents using the C2 Beacon's host node then blocking all traffic would lead to a negative reward. The same applies for the previous example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the C2 Suite using the c2_setup method & capturing the OBS impacts\n",
    "\n",
    "blue_env, c2_server, c2_beacon, client_1, web_server = c2_setup(given_env=blue_env)\n",
    "pre_blue_action_obs, _, _, _, _ = blue_env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code cell below uses the custom blue agent defined at the start of this section to perform a router_acl_add_rule on router 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using CAOS ACTION: router_acl_add_rule & capturing the OBS\n",
    "post_blue_action_obs, _, _, _, _ = blue_env.step(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can see the effects of after another timestep and looking at router 1's ACLs and the OBS differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "router_1: Router = blue_env.game.simulation.network.get_node_by_hostname(\"router_1\")\n",
    "router_1.acl.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can see that the C2 applications are unable to maintain connection - thus being unable to execute correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_env.step(0)\n",
    "\n",
    "# Attempting to install and execute the ransomware script\n",
    "c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)\n",
    "c2_server.send_command(C2Command.RANSOMWARE_LAUNCH, command_options={})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "router_1.acl.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because of the ACL rule the C2 beacon never received the ransomware installation and execute commands from the C2 server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "web_server.software_manager.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "database_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"database_server\")\n",
    "database_server.software_manager.file_system.show(full=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_obs_diffs(pre_blue_action_obs, post_blue_action_obs, blue_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Command and Control** | Configurability \n",
    "\n",
    "This section of the notebook demonstrates the C2 configuration options and their impact on the simulation layer and the game layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The table below is the currently offered C2 Beacon configuration options:\n",
    "\n",
    "|Configuration Option | Option Meaning                                                            |Default Option | Type    | _Optional_ |\n",
    "|---------------------|---------------------------------------------------------------------------|---------------|---------|------------|\n",
    "|c2_server_ip_address | The IP Address of the C2 Server. (The C2 Server must be running)          |_None_         |str (IP) | _No_       |\n",
    "|keep_alive_frequency | How often should the C2 Beacon confirm it's connection in timesteps.      |5              |Int      | _Yes_      |\n",
    "|masquerade_port      | What port should the C2 traffic use? (TCP or UDP)                         |TCP            |Str      | _Yes_      |\n",
    "|masquerade_protocol  | What protocol should the C2 traffic masquerade as? (HTTP, FTP or DNS)     |HTTP           |Str      | _Yes_      |\n",
    "\n",
    "The C2 Server currently does not offer any unique configuration options. The C2 Server aligns itself with the C2 Beacon's configuration options once connection is established."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As demonstrated earlier, red agents can use the ``configure_c2_beacon`` action to configure these settings mid episode through the configuration options:\n",
    "\n",
    "``` YAML\n",
    "...\n",
    "    action: configure_c2_beacon\n",
    "        options:\n",
    "        node_id: 0\n",
    "        config:\n",
    "            c2_server_ip_address: 192.168.10.21\n",
    "            keep_alive_frequency: 10\n",
    "            masquerade_protocol: TCP\n",
    "            masquerade_port: DNS\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | Configurability | C2 Server IP Address\n",
    "\n",
    "As with a majority of client and server based application configurations in primaite, the remote IP of a server must be supplied.\n",
    "\n",
    "In the case of the C2 Beacon, the C2 Server's IP address must be supplied before the C2 beacon will be able to perform any other actions (including ``APPLICATION EXECUTE``).\n",
    "\n",
    "If the network contains multiple C2 Servers then it's also possible to switch to a different C2 server mid-episode which is demonstrated in the below code cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_manipulation_config_path()) as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    # removing all agents & adding the custom agent.\n",
    "    cfg['agents'] = {}\n",
    "    cfg['agents'] = c2_agent_yaml\n",
    "\n",
    "\n",
    "c2_config_env = PrimaiteGymEnv(env_config=cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Installing the C2 Server on both client 1 and client 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "web_server: Server = c2_config_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
    "web_server.software_manager.install(C2Beacon)\n",
    "c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
    "\n",
    "client_1: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
    "client_1.software_manager.install(C2Server)\n",
    "c2_server_1: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
    "c2_server_1.run()\n",
    "\n",
    "client_2: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_2\")\n",
    "client_2.software_manager.install(C2Server)\n",
    "c2_server_2: C2Server = client_2.software_manager.software[\"C2Server\"]\n",
    "c2_server_2.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Configuring the C2 Beacon to establish connection to the C2 Server on client_1 (192.168.10.21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(2) # Agent Action Equivalent to c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
    "env.step(3) # Agent action Equivalent to c2_beacon.establish()\n",
    "c2_beacon.show()\n",
    "c2_server_1.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now reconfiguring the C2 Beacon to establish connection to the C2 Server on client_2 (192.168.10.22)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(9) # Equivalent of to c2_beacon.configure(c2_server_ip_address=\"192.168.10.22\")\n",
    "env.step(3)\n",
    "\n",
    "c2_beacon.show()\n",
    "c2_server_2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After six timesteps the client_1 server will recognise the C2 beacon's previous connection as dead and clear its connections. (This is dependent on the ``Keep Alive Frequency`` setting.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(6):\n",
    "    env.step(0)\n",
    "\n",
    "c2_server_1.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | Configurability | Keep Alive Frequency\n",
    "\n",
    "In order to confirm it's connection the  C2 Beacon will send out a ``Keep Alive`` to the C2 Server and receive a keep alive back. \n",
    "\n",
    "By default, this occurs every 5 timesteps. However, this setting can be configured to be much more infrequent or as frequent as every timestep. \n",
    "\n",
    "The next set of code cells below demonstrate the impact that this setting has on blue agent observation space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_manipulation_config_path()) as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    # removing all agents & adding the custom agent.\n",
    "    cfg['agents'] = {}\n",
    "    cfg['agents'] = custom_blue\n",
    "    cfg['agents'][0]['observation_space']['options']['components'][0]['options']['num_ports'] = 3\n",
    "    cfg['agents'][0]['observation_space']['options']['components'][0]['options']['monitored_traffic'].update({\"tcp\": [\"HTTP\",\"FTP\"]})\n",
    "    cfg['agents'][0]['observation_space']['options']['components'][0]['options']['monitored_traffic'].update({\"udp\": [\"DNS\"]})\n",
    "\n",
    "blue_config_env = PrimaiteGymEnv(env_config=cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Performing the usual c2 setup:\n",
    "blue_config_env, c2_server, c2_beacon, client_1, web_server = c2_setup(given_env=blue_config_env)\n",
    "\n",
    "# Flushing out the OBS impacts from setting up the C2 suite.\n",
    "blue_config_env.step(0)\n",
    "blue_config_env.step(0)\n",
    "\n",
    "# Capturing the 'default' obs (Post C2 installation and configuration):\n",
    "default_obs, _, _, _, _ = blue_config_env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next code cells capture the obs impact of the default Keep Alive Frequency which is 5 timesteps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
    "c2_beacon.establish()\n",
    "c2_beacon.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code cell below executes 10 timesteps and displays the differences between the default and the current timestep.\n",
    "\n",
    "You will notice that the only two timesteps displayed observation space differences. This is due to the C2 Suite confirming their connection through sending ``Keep Alive`` traffic across the network every 5 timesteps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    keep_alive_obs, _, _, _, _ = blue_config_env.step(0)\n",
    "    display_obs_diffs(default_obs, keep_alive_obs, blue_config_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the code cells below configure the C2 Beacon to confirm connection on every timestep via changing the ``keep_alive_frequency`` to ``1``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\", keep_alive_frequency=1)\n",
    "c2_beacon.establish()\n",
    "c2_beacon.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Demonstrating that the observation impacts of the Keep Alive can be seen on every timestep:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparing the OBS of the default frequency to a timestep frequency of 1\n",
    "for i in range(2):\n",
    "    keep_alive_obs, _, _, _, _ = blue_config_env.step(0)\n",
    "    display_obs_diffs(default_obs, keep_alive_obs, blue_config_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, the keep_alive_frequency can also be used to configure the C2 Beacon to confirm connection less frequently. \n",
    "\n",
    "The code cells below demonstrate the impacts of changing the frequency rate to ``7`` timesteps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\", keep_alive_frequency=7)\n",
    "\n",
    "# Comparing the OBS of the default frequency to a timestep frequency of 7\n",
    "for i in range(7):\n",
    "    keep_alive_obs, _, _, _, _ = blue_config_env.step(0)\n",
    "    display_obs_diffs(default_obs, keep_alive_obs, blue_config_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Command and Control** | Configurability | Masquerade Port & Masquerade Protocol\n",
    "\n",
    "The final configurable options are ``Masquerade Port`` & ``Masquerade Protocol``. These options can be used to control the networking IP Protocol and Port the C2 traffic is currently using.\n",
    "\n",
    "In the real world, adversaries take defensive steps to reduce the chance that an installed C2 Beacon is discovered. One of the most commonly used methods is to masquerade C2 traffic as other commonly used networking protocols.\n",
    "\n",
    "In primAITE, red agents can begin to simulate stealth behaviour by configuring C2 traffic to use different protocols mid episode or between episodes.\n",
    "\n",
    "Currently, red agent actions support the following port and protocol options:\n",
    "\n",
    "| Supported Ports  | Supported Protocols |\n",
    "|------------------|---------------------|\n",
    "|``DNS``           | ``UDP``             |\n",
    "|``FTP``           | ``TCP``             |\n",
    "|``HTTP``          |                     |\n",
    "\n",
    "\n",
    "\n",
    "The next set of code cells will demonstrate the impact of this option from a blue agent perspective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_config_env.reset()\n",
    "\n",
    "# Performing the usual c2 setup:\n",
    "blue_config_env, c2_server, c2_beacon, client_1, web_server = c2_setup(given_env=blue_config_env)\n",
    "\n",
    "blue_config_env.step(0)\n",
    "\n",
    "# Capturing the 'default' obs (Post C2 installation and configuration):\n",
    "default_obs, _, _, _, _ = blue_config_env.step(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the C2 suite will masquerade a Web Browser, meaning C2 Traffic will opt to use ``TCP`` and ``HTTP`` (Port 80):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capturing default C2 Traffic\n",
    "for i in range(3):\n",
    "    tcp_c2_obs, _, _, _, _ = blue_config_env.step(0)\n",
    "\n",
    "display_obs_diffs(default_obs, tcp_c2_obs, blue_config_env.game.step_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, C2 Beacon can be configured to use UDP (``Masquerade Protocol``) and we can also configure the C2 Beacon to use a different Port (``Masquerade Port``) for example ``DNS``.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP\n",
    "from primaite.utils.validation.port import PORT_LOOKUP\n",
    "\n",
    "# As we're configuring via the PrimAITE API we need to pass the actual IPProtocol/Port (Agents leverage the simulation via the game layer and thus can pass strings).\n",
    "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\", masquerade_protocol=PROTOCOL_LOOKUP[\"UDP\"],  masquerade_port=PORT_LOOKUP[\"DNS\"])\n",
    "c2_beacon.establish()\n",
    "c2_beacon.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capturing UDP C2 Traffic\n",
    "for i in range(5):\n",
    "    udp_c2_obs, _, _, _, _ = blue_config_env.step(0)\n",
    "\n",
    "display_obs_diffs(tcp_c2_obs, udp_c2_obs, blue_config_env.game.step_counter)"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
