blog_notebooks/cyber/flow_classification/flow_classification_rapids.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cyber Use Case Tutorial: Multiclass Classification on IoT Flow Data with XGBoost"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Goals:\n",
    "- Learn the basics of cyber network data with respect to consumer IoT devices\n",
    "- Load network data into a cuDF\n",
    "- Explore network data and features\n",
    "- Use XGBoost to build a classification model\n",
    "- Evaluate the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get started, we'll make sure the data is available and in the expected location. If you already have the data on your machine, change the `DATA_PATH` location to point to the appropriate location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conn.log\n",
      "lab_mac_labels_cats.csv\n",
      "small_sample.pcap\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p ../../../data/input/unswiot\n",
    "!if [ ! -f ../../../data/input/unswiot/conn.log ]; then tar -xzvf ../../../data/unswiot/unswiotflow.tar.gz -C ../../../data/input/unswiot/; fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify the location of the data files\n",
    "DATA_PATH = \"../../../data/input/unswiot/\"\n",
    "\n",
    "# the sample PCAP file used for explanation \n",
    "DATA_PCAP = DATA_PATH + \"small_sample.pcap\"\n",
    "\n",
    "# the flow connection log (conn.log) file\n",
    "DATA_SOURCE = DATA_PATH + \"conn.log\"\n",
    "\n",
    "# the data label file (matches IP addresses with MAC addresses)\n",
    "DATA_LABELS = DATA_PATH + \"lab_mac_labels_cats.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The Internet of Things and Data at a Massive Scale\n",
    "Gartner estimates there are currently over 8.4 billion Internet of Things (IoT) devices. By 2020, that number is [estimated to surpass 20 billion](https://www.zdnet.com/article/iot-devices-will-outnumber-the-worlds-population-this-year-for-the-first-time/). These types of devices range from consumer devices (e.g., Amazon Echo, smart TVs, smart cameras, door bells) to commercial devices (e.g., building automation systems, keycard entry). All of these devices exhibit behavior on the Internet as they communicate back with their own clouds and user-specified integrations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Types of Network Data\n",
    "The most detailed type of data that is typically collected on a network is full Packet CAPture (PCAP) data. This information is detailed and contains everything about the communication, including: source address, destination address, protocols used, bytes transferred, and even the raw data (e.g., image, audio file, executable). PCAP data is fine-grained, meaning that there is a record for each frame being transmitted. A typical communication is composed of many individual packets/frames.\n",
    "\n",
    "If we aggregate PCAP data so that there is one row of data per communication session, we call that flow level data. A simplified example of this relationship is shown in the figure below.\n",
    "\n",
    "![PCAP_flow_relationship](images/pcap_vs_flow.png \"PCAP vs FLOW\")\n",
    "\n",
    "For this tutorial, we use data from the University of New South Wales. In a lab environment, they [collected nearly three weeks of IoT data from 21 IoT devices](http://149.171.189.1). They also kept a detailed [list of devices by MAC address](http://149.171.189.1/resources/List_Of_Devices.txt), so we have ground-truth with respect to each IoT device's behavior on the network.\n",
    "\n",
    "**Our goal is to utilize the behavior exhibited in the network data to classify IoT devices.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Investigation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first see some of the data. We'll load a PCAP file in using Scapy. If you don't want to or can't install Scapy, feel free to skip this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q scapy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scapy.all import *\n",
    "cap = rdpcap(DATA_PCAP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eth_frame = cap[3]\n",
    "ip_pkt = eth_frame.payload\n",
    "segment = ip_pkt.payload\n",
    "data = segment.payload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eth_frame.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's really a lot of features there. In addition to having multiple layers (which may differ between packets), there are a number of other issues with working directly with PCAP. Often the payload (the `Raw` section above) is encrypted, rendering it useless. The lack of aggregation also makes it difficult to differentiate between packets. What we really care about for this application is what a *session* looks like. In other words, how a Roku interacts with the network is likely quite different than how a Google Home interacts. \n",
    "\n",
    "To save time for the tutorial, all three weeks of PCAP data have already been transformed to flow data, and we can load that in to a typical Pandas dataframe. Due to how the data was created, we have a header row (with column names) as well as a footer row. We've already removed those rows, so nothing to do here.\n",
    "\n",
    "For this application, we used [Zeek](https://www.zeek.org) (formerly known as Bro) to construct the flow data. To include MAC addresses in the conn log, we used the [mac-logging.zeek script](https://github.com/bro/bro/blob/master/scripts/policy/protocols/conn/mac-logging.zeek).\n",
    "\n",
    "If you've skipped installing Scapy, you can pick up here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudf as cd\n",
    "import pandas as pd\n",
    "import nvstrings\n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "pdf = pd.read_csv(DATA_SOURCE, sep='\\t')\n",
    "print(\"==> pdf shape: \",pdf.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at what this new aggregated data looks like, and get a better sense of the columns and their data types. Let's do this the way we're familiar with, using Pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's Pandas, and we could continue the analysis there if we wanted. But what about  [cuDF](https://github.com/rapidsai/cudf)? Let's pivot to that for the majority of this tutorial.\n",
    "\n",
    "One thing cuDF neeeds is for us to specify the data types. We'll write a function to make this easier. As of version 0.6, [strings are supported in cuDF](https://rapidsai.github.io/projects/cudf/en/latest/10min.html?highlight=string#String-Methods). We'll make use of that here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dtypes(fn, delim, floats, strings):\n",
    "    with open(fn, errors='replace') as fp:\n",
    "        header = fp.readline().strip()\n",
    "    \n",
    "    types = []\n",
    "    for col in header.split(delim):\n",
    "        if 'date' in col:\n",
    "            types.append((col, 'date'))\n",
    "        elif col in floats:\n",
    "            types.append((col, 'float64'))\n",
    "        elif col in strings:\n",
    "            types.append((col, 'str'))\n",
    "        else:\n",
    "            types.append((col, 'int64'))\n",
    "\n",
    "    return OrderedDict(types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes_data_processed = get_dtypes(DATA_SOURCE, '\\t', floats=['ts','duration'],\n",
    "                             strings=['uid','id.orig_h','id.resp_h','proto','service',\n",
    "                                      'conn_state','local_orig','local_resp',\n",
    "                                      'history','tunnel_parents','orig_l2_addr',\n",
    "                                      'resp_l2_addr'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "raw_cdf = cd.io.csv.read_csv(DATA_SOURCE, delimiter='\\t', names=list(dtypes_data_processed), \n",
    "                                       dtype=list(dtypes_data_processed.values()), skiprows=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes_data_processed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those data types seem right. Let's see what this data looks like now that it's in cuDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(raw_cdf.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding ground truth labels back to the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll need some labels for our classification task, so we've already prepared a file with those labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes_labels_processed = get_dtypes(DATA_LABELS, ',', floats=[],\n",
    "                             strings=['device','mac','connection','category'])\n",
    "\n",
    "labels_cdf = cd.io.csv.read_csv(DATA_LABELS, delimiter=',', names=list(dtypes_labels_processed), \n",
    "                                       dtype=list(dtypes_labels_processed.values()), skiprows=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(labels_cdf.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes_labels_processed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now perform a series of merges to add the ground truth data (device name, connection, category, and categoryID) back to the dataset. Since each row of netflow has two participants, we'll have to do this twice - once for the originator (source) and once for the responder (destination)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "labels_cdf.columns = ['orig_device','orig_l2_addr','orig_connection','orig_category','orig_category_id']\n",
    "merged_cdf = cd.merge(raw_cdf, labels_cdf, how='left', on='orig_l2_addr')\n",
    "labels_cdf.columns = ['resp_device','resp_l2_addr','resp_connection','resp_category','resp_category_id']\n",
    "merged_cdf = cd.merge(merged_cdf, labels_cdf, how='left')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's reset the `labels_cdf` column names for our own sanity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_cdf.columns = ['device','mac','connection','category','category_id']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's just look at our new dataset to make sure everything's okay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(merged_cdf.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_cdf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exploding the Netflow Data into Originator and Responder Rows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have netflow that has one row per (sessionized) communication between an originator and responder. However, in order to classify an individual device, we need to explode data. Instead of one row that contains both originator and responder, we'll explode to one row for originator information (orig_bytes, orig_pkts, orig_ip_bytes) and one for responder information (resp_bytes, resp_pkts, resp_ip_bytes).\n",
    "\n",
    "The easiest way to do this is to create two new dataframes, rename all of the columns, then `concat` them back together. Just for sanity, we'll also check the new shape of our exploded data frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orig_comms_cdf = merged_cdf[['ts','id.orig_h','id.orig_p','proto','service','duration',\n",
    "                             'orig_bytes','orig_pkts','orig_ip_bytes','orig_device',\n",
    "                             'orig_l2_addr','orig_category','orig_category_id']]\n",
    "orig_comms_cdf.columns = ['ts','ip','port','proto','service','duration','bytes','pkts',\n",
    "                          'ip_bytes','device','mac','category','category_id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp_comms_cdf = merged_cdf[['ts','id.resp_h','id.resp_p','proto','service','duration',\n",
    "                             'resp_bytes','resp_pkts','resp_ip_bytes','resp_device',\n",
    "                             'resp_l2_addr','resp_category','resp_category_id']]\n",
    "resp_comms_cdf.columns = ['ts','ip','port','proto','service','duration','bytes','pkts',\n",
    "                          'ip_bytes','device','mac','category','category_id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exploded_cdf = cd.multi.concat([orig_comms_cdf, resp_comms_cdf])\n",
    "print(\"==> shape (original) =\", merged_cdf.shape)\n",
    "print(\"==> shape =\", exploded_cdf.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're going to need the number of categories (classes) quite a bit, so we'll make a variable for it for easier access. For this tutorial using the data originally presented, we should have 13 categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "num_categories = labels_cdf['category_id'].unique().shape[0]\n",
    "print(\"==> number of IoT categories =\", num_categories)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We currently need to remove null values before we proceed. Although `dropna` doesn't exist in cuDF yet, we can use a workaround to get us there. Also, due to what's available currently, we can't have any nulls in any place in the DF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in exploded_cdf.columns:\n",
    "    print(col, exploded_cdf[col].null_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exploded_cdf['category_id'] = exploded_cdf['category_id'].fillna(-999)\n",
    "exploded_cdf['device'] = exploded_cdf['device'].str.fillna(\"none\")\n",
    "exploded_cdf['category'] = exploded_cdf['category'].str.fillna(\"none\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in exploded_cdf.columns:\n",
    "    print(col, exploded_cdf[col].null_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like all the null values are gone, so now we can proceed. If an IP doesn't have a category ID, we can't use it. So we'll filter those out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exploded_cdf = exploded_cdf[exploded_cdf['category_id'] != -999]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exploded_cdf.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binning the Data and Aggregating the Features\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But wait, there's still more data wrangling to be done! While we've exploded the flows into rows for orig/resp, we may want to bin the data further by time. The rationale is that any single communication may not be an accurate representation of how a device typically reacts in its environment. Imagine the simple case of how a streaming camera typically operates (most of its data will be uploaded from the device to a destination) versus how it operates during a firmware update (most of the data will be pushed down to the device, after which a brief disruption in connectivity will occur).\n",
    "\n",
    "There's a lof ot different time binning we could do. It also would be useful to investigate what the average duration of connection is relative to how many connections per time across various time granularities. With that said, we'll just choose a time bin of 1 hour to begin with. In order to bin, we'll use the following formula:\n",
    "\n",
    "$$\\text{hour_time_bin}=\\left\\lfloor{\\frac{ts}{60*60}}\\right\\rfloor$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "exploded_cdf['hour_time_bin'] = exploded_cdf['ts'].applymap(lambda x: math.floor(x/(60*60))).astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also have to make a choice about how we'll aggregate the binned data. One of the simplest ways is to sum the bytes and packets. There are really two choices for bytes, `bytes` and `ip_bytes`. With Bro, `bytes` is taken from the TCP sequence numbers and is potentially inaccurate, so we select `ip_bytes` instead for both originator and responder. We'll also use the sum of the number of packets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "one_hour_time_bin_cdf = (exploded_cdf[['bytes','pkts','ip_bytes',\n",
    "                                      'mac','category_id',\n",
    "                                      'hour_time_bin']]\n",
    "                            .groupby(['mac','category_id','hour_time_bin'])\n",
    "                            .agg({'bytes':'sum',\n",
    "                                  'pkts':'sum',\n",
    "                                  'ip_bytes':'sum'})\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "one_hour_time_bin_cdf.columns = ['mac', 'category_id', 'hour_time_bin', \n",
    "                                 'bytes', 'pkts', 'ip_bytes']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the Training and Testing Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll take a tradition 70/30 train/test split, and we'll randomly sample into a train and test data frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "cdf_msk = np.random.rand(len(one_hour_time_bin_cdf)) < 0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_cdf = one_hour_time_bin_cdf[cdf_msk]\n",
    "test_cdf = one_hour_time_bin_cdf[~cdf_msk]\n",
    "\n",
    "print(\"==> train length =\",len(train_cdf))\n",
    "print(\"==> test length =\",len(test_cdf))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare the training input (`train_X`), training target (`train_Y`), test input (`test_X`) and test target (`test_Y`) datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = train_cdf[['pkts','ip_bytes']]\n",
    "train_Y = train_cdf[['category_id']]\n",
    "\n",
    "test_X = test_cdf[['pkts','ip_bytes']]\n",
    "test_Y = test_cdf[['category_id']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we just look at the head of both of these datasets (just a quick sanity check)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_X.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_Y.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure XGBoost"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We choose a classification algorithm that utilizes the GPU - [XGBoost](https://xgboost.readthedocs.io/en/latest/). The package provides support for gradient boosted trees and can leverage distributed GPU compute environments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Getting data into a format for XGBoost is really easy. Just make a `DMatrix` for both training and testin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xg_train = xgb.DMatrix(train_X, label=train_Y)\n",
    "xg_test = xgb.DMatrix(test_X, label=test_Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like any good ML package, there's quite a few parameters to set. We're going to start with the softmax objective function. This will let us get a predicted category out of our model. We'll also set other parameters like the maximum depth and number of threads. You can read more about the parameters [here](https://xgboost.readthedocs.io/en/latest/parameter.html). Experiment with them!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param = {}\n",
    "param['objective'] = 'multi:softmax'\n",
    "param['eta'] = 0.1\n",
    "param['max_depth'] = 8\n",
    "param['silent'] = 1\n",
    "param['nthread'] = 4\n",
    "param['num_class'] = num_categories\n",
    "param['max_features'] = 'auto'\n",
    "param['n_gpus'] = 1\n",
    "param['tree_method'] = 'gpu_hist'\n",
    "# param"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "XGBoost allows us to define a watchlist so what we can keep track of performance as the algorithm trains. We'll configure a simple watchlist that is watching `xg_train` and `xg_gest` error rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "watchlist = [(xg_train, 'train'), (xg_test, 'test')]\n",
    "num_round = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training our First XGBoost Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bst = xgb.train(param, xg_train, num_round, watchlist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prediction is also easy (and fast)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = bst.predict(xg_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We might want to get a sense of how our model is by calculating the error rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_cdf = cd.from_pandas(pd.DataFrame(pred, columns=['pred']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_cdf.add_column('category_id',test_Y['category_id'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "error_rate = (pred_cdf[pred_cdf['pred'] != pred_cdf['category_id']]['pred'].count()) / test_Y.shape[0]\n",
    "error_rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's not great, but it's not terrible considering we made quite a few seemingly abritrary decisions in both the feature selection and aggregation phases. Maybe we want to get some more insight into how our model is performing by analyzing the ROC curves for each class, micro average, and macro average. We'll revert back to traditional Python data science tools to do this analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyzing the Model's Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by importing some packages we'll need to perform this analysis. For simplicity in an already large notebook, we'll put them in a single cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sklearn is used to binarize the labels as well as calculate ROC and AUC\n",
    "from sklearn.metrics import roc_curve, auc,recall_score,precision_score\n",
    "from sklearn.preprocessing import label_binarize\n",
    "\n",
    "# scipy is used for interpolating the ROC curves\n",
    "from scipy import interp\n",
    "\n",
    "# our old friend matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "# choose whatever style you want\n",
    "plt.style.use('fivethirtyeight')\n",
    "\n",
    "# cycle is used just to make different colors for the different ROC curves\n",
    "from itertools import cycle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A ROC curve analysis can be trickey for multiclass problems. One way to deal with it is to look at the ROC curve for each class. We'll take some steps to format our data so that it plays nicely with input requirements from sklearn (ah 80/20 rule, we meet again). We also will need to rerun our model with a different objective function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rerunning the Model with the `softprob` Objective Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We used the `softmax` objective function above, but what we really want out of model this time is probabilities that a netflow communication belongs to each of the classes. This is easy enough to do with XGBoost, as we just change the objective function to `softprob`. For simplicity, all of the configuration is in a single cell below rather than spread out. Note the only difference is the objective function change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf_msk = np.random.rand(len(one_hour_time_bin_cdf)) < 0.7\n",
    "\n",
    "train_cdf = one_hour_time_bin_cdf[cdf_msk]\n",
    "test_cdf = one_hour_time_bin_cdf[~cdf_msk]\n",
    "\n",
    "train_X = train_cdf[['pkts','ip_bytes']]\n",
    "train_Y = train_cdf[['category_id']]\n",
    "\n",
    "test_X = test_cdf[['pkts','ip_bytes']]\n",
    "test_Y = test_cdf[['category_id']]\n",
    "\n",
    "xg_train = xgb.DMatrix(train_X, label=train_Y)\n",
    "xg_test = xgb.DMatrix(test_X, label=test_Y)\n",
    "\n",
    "param = {}\n",
    "param['objective'] = 'multi:softprob'\n",
    "param['eta'] = 0.1\n",
    "param['max_depth'] = 8\n",
    "param['silent'] = 1\n",
    "param['nthread'] = 4\n",
    "param['num_class'] = num_categories\n",
    "param['n_gpus'] = 1\n",
    "param['tree_method'] = 'gpu_hist'\n",
    "\n",
    "watchlist = [(xg_train, 'train'), (xg_test, 'test')]\n",
    "num_round = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bst = xgb.train(param, xg_train, num_round, watchlist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so we have our new model. We now take some steps to make sure the data is in a format that makes sklearn happy. First we'll use the `predict` function to compute the probabilities. To extend `roc_curve` to multiclass, we'll also need to binarize the labels. Let's keep our sanity by also making sure the lengths match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(bst.predict(xg_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probs = bst.predict(xg_test).reshape(test_Y.shape[0],param['num_class'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For now, we need to convert the `test_Y` cuDF to an array. The most straightforward way to do that is to go through Pandas. It also lets us show off how nicely we can convert to Pandas, should the need arise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_Y_binarize = label_binarize(test_Y.to_pandas()['category_id'].values, classes=np.arange(param['num_class']))\n",
    "\n",
    "print(\"==> length of probs =\",len(probs))\n",
    "print(\"==> length of test_Y_binarize =\", len(test_Y_binarize))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some more housekeeping. We'll create Python dictionaries to hold FPR ([false positive rate](https://en.wikipedia.org/wiki/False_positive_rate)), TPR ([true positive rate](https://en.wikipedia.org/wiki/Sensitivity_and_specificity)), and AUC ([area under the curve](https://en.wikipedia.org/wiki/Receiver_operating_characteristic#Area_under_the_curve)) values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fpr = dict()\n",
    "tpr = dict()\n",
    "roc_auc = dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each of our classes, we'll computer FPR, TPR, and AUC. We're also compute the [micro and macro averages](http://rushdishams.blogspot.com/2011/08/micro-and-macro-average-of-precision.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"==> number of classes =\", num_categories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate FPR, TPR, and ROC AUC for every class\n",
    "for i in range(num_categories):\n",
    "    fpr[i], tpr[i], _ = roc_curve(test_Y_binarize[:, i], probs[:, i])\n",
    "    roc_auc[i] = auc(fpr[i], tpr[i])\n",
    "\n",
    "# calculate the micro average FPR, TPR, and ROC AUC (we'll calculate the macro average below)\n",
    "fpr[\"micro\"], tpr[\"micro\"], _ = roc_curve(test_Y_binarize.ravel(), probs.ravel())\n",
    "roc_auc[\"micro\"] = auc(fpr[\"micro\"], tpr[\"micro\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the ROC Curves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phew! Lots of code below, but it's fairly straightforward and [adapted from an example in the scikit-learn documentation](http://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html#multiclass-settings). Before we plot though, we'll create a simple category lookup dictionary so we can label the classes with their actual names (not their category IDs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_pdf = labels_cdf.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "category_lookup = labels_pdf[['category','category_id']].drop_duplicates().set_index('category_id').T.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregate all of the false positive rates across all classes\n",
    "all_fpr = np.unique(np.concatenate([fpr[i] for i in range(num_categories)]))\n",
    "\n",
    "# interpolate all of the ROC curves\n",
    "mean_tpr = np.zeros_like(all_fpr)\n",
    "for i in range(param['num_class']):\n",
    "    mean_tpr += interp(all_fpr, fpr[i], tpr[i])\n",
    "\n",
    "# average the TPR\n",
    "mean_tpr /= num_categories\n",
    "\n",
    "# compute the macro average FPR, TPR, and ROC AUC\n",
    "fpr['macro'] = all_fpr\n",
    "tpr['macro'] = mean_tpr\n",
    "roc_auc['macro'] = auc(fpr['macro'], tpr['macro'])\n",
    "\n",
    "# plot all of the ROC curves on a single plot (for comparison)\n",
    "plt.figure(figsize=(9,9))\n",
    "plt.plot(fpr['micro'], tpr['micro'],\n",
    "         label=\"micro-average ROC curve (area = {0:0.2f})\"\n",
    "               \"\".format(roc_auc['micro']),\n",
    "         color='deeppink', linestyle=':', linewidth=4)\n",
    "\n",
    "plt.plot(fpr['macro'], tpr['macro'],\n",
    "         label=\"macro-average ROC curve (area = {0:0.2f})\"\n",
    "               \"\".format(roc_auc['macro']),\n",
    "         color='navy', linestyle=':', linewidth=4)\n",
    "\n",
    "num_colors = param['num_class']\n",
    "cm = plt.get_cmap('gist_rainbow')\n",
    "\n",
    "colors = cycle([cm(1.*i/num_colors) for i in range(num_colors)])\n",
    "\n",
    "lw = 2\n",
    "for i, color in zip(range(param['num_class']), colors):\n",
    "    plt.plot(fpr[i], tpr[i], color=color, lw=lw,\n",
    "             label=\"ROC curve for \"+category_lookup[i]['category']+\" class (area = {1:0.2f})\"\n",
    "             \"\".format(i, roc_auc[i]))\n",
    "\n",
    "plt.plot([0, 1], [0, 1], 'k--', lw=lw)\n",
    "plt.xlim([0.0, 1.0])\n",
    "plt.ylim([0.0, 1.05])\n",
    "plt.xlabel(\"False Positive Rate\", fontsize=12)\n",
    "plt.ylabel(\"True Positive Rate\", fontsize=12)\n",
    "plt.title(\"ROC Curves for IoT Device Categories\")\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's not a *terrible* plot, but it gets a little messy. We can also plot each class as its own subplot.\n",
    "\n",
    "First we make a few variables so we can control the layout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_subplots = num_categories\n",
    "plot_grid_cols = 3\n",
    "plot_grid_rows = total_subplots // plot_grid_cols\n",
    "plot_grid_rows += total_subplots % plot_grid_cols\n",
    "\n",
    "position_index = range(1, total_subplots+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we make the grid of plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "fig, axs = plt.subplots(plot_grid_rows, plot_grid_cols, sharex=True, sharey=True, figsize=(15,15))\n",
    "\n",
    "lw = 2\n",
    "\n",
    "plt_num = 0\n",
    "for row in range(plot_grid_rows):\n",
    "    for col in range(plot_grid_cols):\n",
    "        if(plt_num <= 12):\n",
    "            axs[row,col].plot(fpr[plt_num], tpr[plt_num], lw=lw)\n",
    "            axs[row,col].set_title(category_lookup[plt_num]['category']+' Devices ROC Curve', fontsize=14)\n",
    "            axs[row,col].text(0.7, 0.1,\"AUC = {:.4f}\".format(roc_auc[plt_num]), size=11)\n",
    "        elif(plt_num == 13):\n",
    "            axs[row,col].plot(fpr['micro'], tpr['micro'], lw=lw)\n",
    "            axs[row,col].set_title(\"Micro Average ROC Curve\", fontsize=14)\n",
    "            axs[row,col].text(0.7, 0.1,\"AUC = {:.4f}\".format(roc_auc['micro']), size=12)\n",
    "        elif(plt_num == 14):\n",
    "            axs[row,col].plot(fpr['macro'], tpr['macro'], lw=lw)\n",
    "            axs[row,col].set_title(\"Macro Average ROC Curve\", fontsize=14)\n",
    "            axs[row,col].text(0.7, 0.1,\"AUC = {:.4f}\".format(roc_auc['macro']), size=12)\n",
    "        axs[row,col].set_xlabel('False Positive Rate', fontsize=10)\n",
    "        axs[row,col].set_ylabel('True Positive Rate', fontsize=10)\n",
    "        plt_num += 1\n",
    "            \n",
    "plt.xlim([-0.01, 1.0])\n",
    "plt.ylim([0.0, 1.05])\n",
    "plt.subplots_adjust(wspace=0.2, hspace=0.4)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we've shown, it's possible to get fairly decent multiclass classification results for IoT data using only basic features (bytes and packets) when aggregated. This isn't surprising, based on the fact that we used expert knowledge to assign category labels. In addition, the majority of the time, IoT devices are in a \"steady state\" (idle), and are not heavily influenced by human interaction. This lets us take larger samples (e.g., aggregate to longer time bins) while still maintaining decent classification performance. It should also be noted that this is a very clean dataset. The traffic is mainly IoT traffic (e.g., little traditional compute traffic), and there are no intentional abnormal activities injected (e.g., red teaming).\n",
    "\n",
    "We used Bro data, but it's also possible to use the raw PCAP data as input for classification. The preprocessing steps are more arduous than for flow data though. It'd be a great exercise..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More to Explore: Possible Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### (1) It may be useful to investigate other time binnings. Can you build another model that uses data binned to a different granularity (e.g., 5 minutes)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your work here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### (2) We used the `sum` of bytes and packets for a device when aggregated to the hour. What about other ways to handle these quantitative features (e.g., average)? Would that improve the classification results?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your work here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### (3) We selected specific parameters for XGBoost. These could probably use a bit more thought. You can [read more about the parameters](https://xgboost.readthedocs.io/en/latest/parameter.html) and try adjusting them on our previous dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a reminder about our parameters\n",
    "print(param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your work here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### (4) There are additional features in the netflow data that we didn't use. Some other quantitative fields (e.g., duration) and categorical fields (e.g., protocol, service, ports) may be useful for classification. Build another XGBoost model using some/all of these fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your work here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Nadji, Y., \"Passive DNS-based Device Identification\", *NANOG 67*, https://www.nanog.org/sites/default/files/Nadji.pdf.\n",
    "1. Shams, R., \"Micro- and Macro-average of Precision, Recall, and F-Score\", http://rushdishams.blogspot.com/2011/08/micro-and-macro-average-of-precision.html.\n",
    "1. Sivanathan, A. et al., \"Characterizing and Classifying IoT Traffic in Smart Cities and Campuses\", *2017 IEEE Conference on Computer Communications Workshops*, May 2017, http://www2.eet.unsw.edu.au/~vijay/pubs/conf/17infocom.pdf.\n",
    "1. University of New South Wales Internet of Things Network Traffic Data Collection, http://149.171.189.1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}