{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Manipulation Scenario\n",
    "\n",
    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario\n",
    "\n",
    "The network consists of an office subnet and a server subnet. Clients in the office access a website which fetches data from a database. Occasionally, admins need to access the database directly from the clients.\n",
    "\n",
    "![UC2 Network](./_package_data/uc2_network.png)\n",
    "\n",
    "_(click image to enlarge)_\n",
    "\n",
    "The red agent deletes the contents of the database. When this happens, the web app cannot fetch data and users navigating to the website get a 404 error.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network\n",
    "\n",
    "- The web server has:\n",
    "    - a web service that replies to user HTTP requests\n",
    "    - a database client that fetches data for the web service\n",
    "- The database server has:\n",
    "    - a POSTGRES database service\n",
    "    - a database file which is accessed by the database service\n",
    "    - FTP client used for backing up the data to the backup_server\n",
    "- The backup server has:\n",
    "    - a copy of the database file in a known good state\n",
    "    - FTP server that can send the backed up file back to the database server\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Green agent\n",
    "\n",
    "There are green agents logged onto client 1 and client 2. They use the web browser to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available.\n",
    "\n",
    "Sometimes, the green agents send a request directly to the database to check that it is reachable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Red agent\n",
    "\n",
    "At the start of every episode, the red agent randomly chooses either client 1 or client 2 to login to. It waits a bit then sends a DELETE query to the database from its chosen client. If the delete is successful, the database file is flagged as compromised to signal that data is not available.\n",
    "\n",
    "![uc2_attack](./_package_data/uc2_attack.png)\n",
    "\n",
    "_(click image to enlarge)_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Blue agent\n",
    "\n",
    "The blue agent can view the entire network, but the health statuses of components are not updated until a scan is performed. The blue agent should restore the database file from backup after it was compromised. It can also prevent further attacks by blocking the red agent client from sending the malicious SQL query to the database server. This can be done by implementing an ACL rule on the router.\n",
    "\n",
    "However, these rules will also impact greens' ability to check the database connection. The blue agent should only block the infected client, it should let the other client connect freely. Once the attack has begun, automated traffic monitoring will detect it as suspicious network traffic. The blue agent's observation space will show this as an increase in the number of malicious network events (NMNE) on one of the network interfaces. To achieve optimal reward, the agent should only block the client which has the non-zero outbound NMNE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reinforcement learning details"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scripted agents:\n",
    "### Red\n",
    "The red agent sits on a client and uses an application called DataManipulationBot whose sole purpose is to send a DELETE query to the database.\n",
    "The red agent can choose one of two action each timestep:\n",
    "1. do nothing\n",
    "2. execute the data manipulation application\n",
    "The schedule for selecting when to execute the application is controlled by three parameters:\n",
    "- start time\n",
    "- frequency\n",
    "- variance\n",
    "\n",
    "Attacks start at a random timestep between (start_time - variance) and (start_time + variance). After each attack, another is attempted after a random delay between (frequency - variance) and (frequency + variance) timesteps.\n",
    "\n",
    "The data manipulation app itself has an element of randomness because the attack has a probability of success. The default is 0.8 to succeed with the port scan step and 0.8 to succeed with the attack itself.\n",
    "Upon a successful attack, the database file becomes corrupted which incurs a negative reward for the RL defender.\n",
    "\n",
    "The red agent does not use information about the state of the network to decide its action.\n",
    "\n",
    "### Green\n",
    "The green agents use the web browser application to send requests to the web server. The schedule of each green agent is currently random, it will do nothing 30% of the time, send a web request 60% of the time, and send a db status check 10% of the time.\n",
    "\n",
    "When a green agent is blocked from accessing the data through the webpage, this incurs a negative reward to the RL defender.\n",
    "\n",
    "Also, when the green agent is blocked from checking the database status, it causes a small negative reward."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observation Space\n",
    "\n",
    "The blue agent's observation space is structured as nested dictionary with the following information:\n",
    "```\n",
    "\n",
    "- NODES\n",
    "    - <node_id 1-7>\n",
    "        - SERVICES\n",
    "            - <service_id 1-1>\n",
    "                - operating_status\n",
    "                - health_status\n",
    "        - FOLDERS\n",
    "            - <folder_id 1-1>\n",
    "                - health_status\n",
    "                - FILES\n",
    "                    - <file_id 1-1>\n",
    "                    - health_status\n",
    "        - NETWORK_INTERFACES\n",
    "            - <nic_id 1-2>\n",
    "                - nic_status\n",
    "                - nmne\n",
    "                    - inbound\n",
    "                    - outbound\n",
    "        - operating_status\n",
    "- LINKS\n",
    "    - <link_id 1-10>\n",
    "        - PROTOCOLS\n",
    "            - ALL\n",
    "                - load\n",
    "- ACL\n",
    "    - <rule_number 1-10>\n",
    "        - position\n",
    "        - permission\n",
    "        - source_node_id\n",
    "        - source_port\n",
    "        - dest_node_id\n",
    "        - dest_port\n",
    "        - protocol\n",
    "- ICS\n",
    "```\n",
    "\n",
    "### Mappings\n",
    "\n",
    "The dict keys for `node_id` are in the following order:\n",
    "\n",
    "| node_id | node name        |\n",
    "|---------|------------------|\n",
    "| 0       | domain_controller|\n",
    "| 1       | web_server       |\n",
    "| 2       | database_server  |\n",
    "| 3       | backup_server    |\n",
    "| 4       | security_suite   |\n",
    "| 5       | client_1         |\n",
    "| 6       | client_2         |\n",
    "\n",
    "Service 1 on node 2 (web_server) corresponds to the Web Server service. Other services are only there for padding to ensure that each node's observation space has the same shape. They are filled with zeroes.\n",
    "\n",
    "Folder 1 on node 3 corresponds to the database folder. File 1 in that folder corresponds to the database storage file. Other files and folders are only there for padding to ensure that each node's observation space has the same shape. They are filled with zeroes.\n",
    "\n",
    "The dict keys for `link_id` are in the following order:\n",
    "\n",
    "| link_id | endpoint_a       | endpoint_b        |\n",
    "|---------|------------------|-------------------|\n",
    "| 1       | router_1         | switch_1          |\n",
    "| 2       | router_1         | switch_2          |\n",
    "| 3       | switch_1         | domain_controller |\n",
    "| 4       | switch_1         | web_server        |\n",
    "| 5       | switch_1         | database_server   |\n",
    "| 6       | switch_1         | backup_server     |\n",
    "| 7       | switch_1         | security_suite    |\n",
    "| 8       | switch_2         | client_1          |\n",
    "| 9       | switch_2         | client_2          |\n",
    "| 10      | switch_2         | security_suite    |\n",
    "\n",
    "\n",
    "The ACL rules in the observation space appear in the same order that they do in the actual ACL. Though, only the first 10 rules are shown, there are default rules lower down that cannot be changed by the agent. The extra rules just allow the network to function normally, by allowing pings, ARP traffic, etc.\n",
    "\n",
    "Most nodes have only 1 network_interface, so the observation for those is placed at NIC index 1 in the observation space. Only the security suite has 2 NICs, the second NIC in the observation space is the one that connects the security suite with swtich_2.\n",
    "\n",
    "The meaning of the services' operating_state is:\n",
    "\n",
    "| operating_state | label      |\n",
    "|-----------------|------------|\n",
    "| 0               | UNUSED     |\n",
    "| 1               | RUNNING    |\n",
    "| 2               | STOPPED    |\n",
    "| 3               | PAUSED     |\n",
    "| 4               | DISABLED   |\n",
    "| 5               | INSTALLING |\n",
    "| 6               | RESTARTING |\n",
    "\n",
    "The meaning of the services' health_state is:\n",
    "\n",
    "| health_state | label       |\n",
    "|--------------|-------------|\n",
    "| 0            | UNUSED      |\n",
    "| 1            | GOOD        |\n",
    "| 2            | FIXING      |\n",
    "| 3            | COMPROMISED |\n",
    "| 4            | OVERWHELMED |\n",
    "\n",
    "\n",
    "The meaning of the files' and folders' health_state is:\n",
    "\n",
    "| health_state | label       |\n",
    "|--------------|-------------|\n",
    "| 0            | UNUSED      |\n",
    "| 1            | GOOD        |\n",
    "| 2            | COMPROMISED |\n",
    "| 3            | CORRUPT     |\n",
    "| 4            | RESTORING   |\n",
    "| 5            | REPAIRING   |\n",
    "\n",
    "\n",
    "The meaning of the NICs' operating_status is:\n",
    "\n",
    "| operating_status | label    |\n",
    "|------------------|----------|\n",
    "| 0                | UNUSED   |\n",
    "| 1                | ENABLED  |\n",
    "| 2                | DISABLED |\n",
    "\n",
    "\n",
    "NMNE (number of malicious network events) means, for inbound or outbound traffic, means:\n",
    "\n",
    "| value | NMNEs          |\n",
    "|-------|----------------|\n",
    "| 0     | None           |\n",
    "| 1     | 1 - 5          |\n",
    "| 2     | 6 - 10         |\n",
    "| 3     | More than 10   |\n",
    "\n",
    "\n",
    "Link load has the following meaning:\n",
    "\n",
    "| load | percent utilisation |\n",
    "|------|---------------------|\n",
    "| 0    | exactly 0%          |\n",
    "| 1    | 0-11%               |\n",
    "| 2    | 11-22%              |\n",
    "| 3    | 22-33%              |\n",
    "| 4    | 33-44%              |\n",
    "| 5    | 44-55%              |\n",
    "| 6    | 55-66%              |\n",
    "| 7    | 66-77%              |\n",
    "| 8    | 77-88%              |\n",
    "| 9    | 88-99%              |\n",
    "| 10   | exactly 100%        |\n",
    "\n",
    "\n",
    "ACL permission has the following meaning:\n",
    "\n",
    "| permission | label  |\n",
    "|------------|--------|\n",
    "| 0          | UNUSED |\n",
    "| 1          | ALLOW  |\n",
    "| 2          | DENY   |\n",
    "\n",
    "\n",
    "ACL source / destination node ids actually correspond to IP addresses (since ACLs work with IP addresses)\n",
    "\n",
    "| source / dest node id | ip_address     | label                   |\n",
    "|-----------------------|----------------|-------------------------|\n",
    "| 0                     |                | UNUSED                  |\n",
    "| 1                     |                | ALL addresses           |\n",
    "| 2                     | 192.168.1.10   | domain_controller       |\n",
    "| 3                     | 192.168.1.12   | web_server              |\n",
    "| 4                     | 192.168.1.14   | database_server         |\n",
    "| 5                     | 192.168.1.16   | backup_server           |\n",
    "| 6                     | 192.168.1.110  | security_suite (eth-1) |\n",
    "| 7                     | 192.168.10.21  | client_1                |\n",
    "| 8                     | 192.168.10.22  | client_2                |\n",
    "| 9                     | 192.168.10.110 | security_suite (eth-2) |\n",
    "\n",
    "\n",
    "ACL source / destination port ids have the following encoding:\n",
    "\n",
    "| port id | port number | port use        |\n",
    "|---------|-------------|-----------------|\n",
    "| 0       |             | UNUSED          |\n",
    "| 1       |             | ALL             |\n",
    "| 2       | 219         | ARP             |\n",
    "| 3       | 53          | DNS             |\n",
    "| 4       | 80          | HTTP            |\n",
    "| 5       | 5432        | POSTGRES_SERVER |\n",
    "\n",
    "\n",
    "ACL protocol ids have the following encoding:\n",
    "\n",
    "| protocol id | label |\n",
    "|-------------|-------|\n",
    "| 0           | UNUSED|\n",
    "| 1           | ALL   |\n",
    "| 2           | ICMP  |\n",
    "| 3           | TCP   |\n",
    "| 4           | UDP   |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Action Space\n",
    "\n",
    "The blue agent chooses from a list of 54 pre-defined actions. The full list is defined in the `action_map` in the config. The most important ones are explained here:\n",
    "\n",
    "- `0`: Do nothing\n",
    "- `1`: Scan the web service - this refreshes the health status in the observation space\n",
    "- `9`: Scan the database file - this refreshes the health status of the database file\n",
    "- `13`: Patch the database service - This triggers the database to restore data from the backup server\n",
    "- `39`: Shut down client 1\n",
    "- `40`: Start up client 1\n",
    "- `46`: Block outgoing traffic from client 1\n",
    "- `47`: Block outgoing traffic from client 2\n",
    "- `50`: Block TCP traffic from client 1 to the database node\n",
    "- `51`: Block TCP traffic from client 2 to the database node\n",
    "- `52-61`: Remove ACL rules 1-10\n",
    "- `66`: Disconnect client 1 from the network\n",
    "- `67`: Reconnect client 1 to the network\n",
    "- `68`: Disconnect client 2 from the network\n",
    "- `69`: Reconnect client 2 to the network\n",
    "\n",
    "The other actions will either have no effect or will negatively impact the network, so the blue agent should avoid taking them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Function\n",
    "\n",
    "The blue agent's reward is calculated using these measures:\n",
    "1. Whether the database file is in a good state (+1 for good, -1 for corrupted, 0 for any other state)\n",
    "2. Whether each green agents' most recent webpage request was successful (+1 for a `200` return code, -1 for a `404` return code and 0 otherwise).\n",
    "3. Whether each green agents' most recent DB status check was successful (+1 for a successful connection, -1 for no connection).\n",
    "\n",
    "The file status reward and the two green-agent-related rewards are averaged to get a total step reward.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demonstration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, load the required modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!primaite setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "from primaite.config.load import data_manipulation_config_path\n",
    "from primaite.session.environment import PrimaiteGymEnv\n",
    "from primaite.game.agent.interface import AgentHistoryItem\n",
    "import yaml\n",
    "from pprint import pprint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the environment. \n",
    "We will also disable the agent observation flattening.\n",
    "\n",
    "This cell will print the observation when the network is healthy. You should be able to verify Node file and service statuses against the description above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the env\n",
    "with open(data_manipulation_config_path(), 'r') as f:\n",
    "    cfg = yaml.safe_load(f)\n",
    "    # set success probability to 1.0 to avoid rerunning cells.\n",
    "    cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
    "    cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
    "    cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
    "    cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
    "    # don't flatten observations so that we can see what is going on\n",
    "    cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n",
    "\n",
    "env = PrimaiteGymEnv(env_config = cfg)\n",
    "obs, info = env.reset()\n",
    "print('env created successfully')\n",
    "pprint(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The red agent will start attacking at some point between step 20 and 30. When this happens, the reward will drop immediately, then drop to -0.8 when green agents try to access the webpage."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(35):\n",
    "    obs, reward, terminated, truncated, info = env.step(0)\n",
    "    print(f\"step: {env.game.step_counter}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the reward is -0.8, let's have a look at blue agent's observation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(obs['NODES'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The true statuses of the database file and webapp are not updated. The blue agent needs to perform a scan to see that they have degraded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs, reward, terminated, truncated, info = env.step(9)  # scan database file\n",
    "obs, reward, terminated, truncated, info = env.step(1)  # scan webapp service\n",
    "\n",
    "pprint(obs['NODES'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now service 1 on HOST1 has `health_status = 3`, indicating that the webapp is compromised.\n",
    "File 1 in folder 1 on HOST2 has `health_status = 2`, indicating that the database file is compromised."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, the NMNE outbound of either client 1 (node 6) or client 2 (node 7) has increased from 0 to 1. This tells us which client is being used by the red agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The blue agent can now patch the database to restore the file to a good health status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs, reward, terminated, truncated, info = env.step(13)  # patch the database\n",
    "print(f\"step: {env.game.step_counter}\")\n",
    "print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
    "print(f\"Green action: {info['agent_actions']['client_1_green_user'].action}\" )\n",
    "print(f\"Green action: {info['agent_actions']['client_2_green_user'].action}\" )\n",
    "print(f\"Blue reward:{reward}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fixing takes two steps, so the reward hasn't changed yet. Let's do nothing for another timestep, the reward should improve.\n",
    "\n",
    "The reward will increase slightly as soon as the file finishes restoring. Then, the reward will increase to 0.9 when both green agents make successful requests.\n",
    "\n",
    "Run the following cell until the green action is `node_application_execute` for application 0, then the reward should increase. If you run it enough times, another red attack will happen and the reward will drop again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs, reward, terminated, truncated, info = env.step(0)  # do nothing\n",
    "print(f\"step: {env.game.step_counter}\")\n",
    "print(f\"Red action: {info['agent_actions']['data_manipulation_attacker'].action}\" )\n",
    "print(f\"Green action: {info['agent_actions']['client_2_green_user']}\" )\n",
    "print(f\"Green action: {info['agent_actions']['client_1_green_user']}\" )\n",
    "print(f\"Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The blue agent can prevent attacks by implementing an ACL rule to stop client_1 or client_2 from sending POSTGRES traffic to the database. (Let's also patch the database file to get the reward back up.)\n",
    "\n",
    "Let's block both clients from communicating directly with the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(13)  # Patch the database\n",
    "print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
    "\n",
    "env.step(50)  # Block client 1\n",
    "print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
    "\n",
    "env.step(51)  # Block client 2\n",
    "print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
    "\n",
    "while abs(reward - 0.8) > 1e-5:\n",
    "    obs, reward, terminated, truncated, info = env.step(0)  # do nothing\n",
    "    print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
    "    if env.game.step_counter > 10000:\n",
    "        break # make sure there's no infinite loop if something went wrong"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, even though the red agent executes an attack, the reward will stay at 0.8."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also have a look at the ACL observation to verify our new ACL rule at positions 5 and 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs['NODES']['ROUTER0']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can slightly increase the reward by unblocking the client which isn't being used by the attacker. If node 6 has outbound NMNEs, let's unblock client 2, and if node 7 has outbound NMNEs, let's unblock client 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.step(58) # Remove the ACL rule that blocks client 1\n",
    "env.step(57) # Remove the ACL rule that blocks client 2\n",
    "\n",
    "tries = 0\n",
    "while True:\n",
    "    tries += 1\n",
    "    obs, reward, terminated, truncated, info = env.step(0)\n",
    "\n",
    "    if obs['NODES']['HOST5']['NICS'][1]['NMNE']['outbound'] == 1:\n",
    "        # client 1 has NMNEs, let's block it\n",
    "        obs, reward, terminated, truncated, info = env.step(50) # block client 1\n",
    "        print(\"blocking client 1\")\n",
    "        break\n",
    "    elif obs['NODES']['HOST6']['NICS'][1]['NMNE']['outbound'] == 1:\n",
    "        # client 2 has NMNEs, so let's block it\n",
    "        obs, reward, terminated, truncated, info = env.step(51) # block client 2\n",
    "        print(\"blocking client 2\")\n",
    "        break\n",
    "    if tries>100:\n",
    "        print(\"Error: NMNE never increased\")\n",
    "        break\n",
    "\n",
    "env.step(13)  # Patch the database\n",
    "print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the reward will eventually increase to 0.9, even after red agent attempts subsequent attacks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(40):\n",
    "    obs, reward, terminated, truncated, info = env.step(0)  # do nothing\n",
    "    print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reset the environment, you can rerun the other cells to verify that the attack works the same every episode. (except the red agent will move between `client_1` and `client_2`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()"
   ]
  }
 ],
 "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
}
