model_comparison.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison of various classifier types for detecting DNS exfiltration\n",
    "\n",
    "This notebook demonstrates the performance of different machine learning models\n",
    "for the task of classifying DNS requests as either legitimate or exfiltration.\n",
    "\n",
    "Prior to running the notebook, you need to download the csv's with the data, \n",
    "available [here](https://data.mendeley.com/datasets/c4n7fckkz3).\n",
    "\n",
    "See the docstrings in `exfiltration_classifier.py` and `exfiltration_dataset.py`\n",
    "for details on the usage of the `ExfiltrationClassifier` and `ExfiltrationDataset`\n",
    "classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from exfiltration_dataset import ExfiltrationDataset\n",
    "from exfiltration_classifier import ExfiltrationClassifier\n",
    "import pprint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the paths to the csv's here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = Path('..')\n",
    "csv_path_modified = root / 'data' / 'dataset_modified.csv'\n",
    "csv_path = root / 'data' / 'dataset.csv'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose the type of features to train the classifiers on:\n",
    "* `'individual'` features are derived from single requests\n",
    "* `'aggregated'` features are derived from sequences of 10 consecutive requests with the same user id and domain\n",
    "* `'all'` means use both individual and aggregated features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_type = 'all' "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We print the following metrics, evaluated on the original examples, \n",
    "modified exfiltrations, and both original and modified examples:\n",
    "* accuracy, as a ratio of correctly classified examples\n",
    "* F1-score\n",
    "    * for legitimate examples (class `False`)\n",
    "    * for exfiltrations (class `True`)\n",
    "    * macro average of the two"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(clf, pp):\n",
    "    report = clf.evaluate('test', 'unmodified')\n",
    "    print('Evaluation on original examples:')\n",
    "    pp.pprint(report)\n",
    "    report = clf.evaluate('test', 'modified')\n",
    "    print('Evaluation on modified exfiltrations:')\n",
    "    pp.pprint(report)\n",
    "    report = clf.evaluate('test', 'both')\n",
    "    print('Evaluation on all examples:')\n",
    "    pp.pprint(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The csv with the original examples contains more than 30M rows.\n",
    "Training the classifiers on the entire dataset could take a long time.\n",
    "You can specify the size of the train, validation and tests sets for both\n",
    "the original examples and the modified exfiltrations. For example, specifying\n",
    "`train_size=0.1`, `val_size=0.1` and `test_size=0.1` means each of the splits\n",
    "will have around 3M examples. The splits is stratified so that each set has\n",
    "approximatly the same ratio of legitimate requests and DNS exfiltrations.\n",
    "Fix the `random_state` parameter for reproducibility.\n",
    "\n",
    "Please note that the csv's are read one line at a time, and this could take some\n",
    "time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading original examples:\n",
      "Reading row   350000 of   350741, 99.79% completed\n",
      "Reading features of modified examples:\n",
      "train: X.shape = (367661, 17), y.shape=(367661,)\n",
      "Reading original examples:\n",
      "Reading row   350000 of   350741, 99.79% completed\n",
      "Reading features of modified examples:\n",
      "test: X.shape = (376121, 17), y.shape=(376121,)\n"
     ]
    }
   ],
   "source": [
    "# Get objects that assign csv rows across splits\n",
    "splitter, splitter_modified = ExfiltrationDataset.create_splits(\n",
    "    csv_path, csv_path_modified, val_size=0, test_size=0.01,\n",
    "    val_size_modified=0, test_size_modified=0.6,\n",
    "    train_size=0.01, train_size_modified=0.4,\n",
    "    random_state=0\n",
    ")\n",
    "\n",
    "# Make datasets\n",
    "ds = {}\n",
    "for split in ['train', 'test']:\n",
    "    ds[split] = ExfiltrationDataset(\n",
    "        feature_type, csv_path, csv_path_modified, \n",
    "        split, splitter, splitter_modified, verbose=True\n",
    "    )\n",
    "    print(f'{split}: X.shape = {ds[split].X.shape}, y.shape={ds[split].y.shape}')\n",
    "\n",
    "# Standardization\n",
    "ds['train'].standardize()\n",
    "ds['test'].standardize(ds['train'].scaler)\n",
    "\n",
    "pp = pprint.PrettyPrinter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Place the types of classifiers you which to train in the `classifier_types` list.\n",
    "Allowed options are:\n",
    "* `'logistic regression'`\n",
    "* `'svm'` (support vector machine with a Gaussian kernel)\n",
    "* `'sgd'` (support vector machine with a linear kernel trained with stochastic gradient descent)\n",
    "* `'naive Bayes`'\n",
    "* `'decision tree'`\n",
    "* `'random forest'`\n",
    "* `'extra trees'`\n",
    "* `'adaboost'`\n",
    "* `'hgb'` (histogram-based gradient boosting)\n",
    "* `'mlp'` (multi-layer perceptron)\n",
    "* `'xgb gpu'` (XGBoost trained on a GPU)\n",
    "* `'xgb cpu'` (XGBoost trained on a CPU)\n",
    "\n",
    "To train only on original examples, and not on instances of exfiltrations\n",
    "obtained with the modified exfiltrator, choose `include_modified=False` in the\n",
    "call to the `train` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "naive Bayes\n",
      "Evaluation on original examples:\n",
      "{'accuracy': 0.9734533459162173,\n",
      " 'f1-score': {'False': 0.9864872905646138,\n",
      "              'True': 0.2507443469864006,\n",
      "              'macro avg': 0.6186158187755072}}\n",
      "Evaluation on modified exfiltrations:\n",
      "{'accuracy': 0.9630023640661939}\n",
      "Evaluation on all examples:\n",
      "{'accuracy': 0.9727481315853143,\n",
      " 'f1-score': {'False': 0.9851447983605655,\n",
      "              'True': 0.8353360750546203,\n",
      "              'macro avg': 0.9102404367075929}}\n"
     ]
    }
   ],
   "source": [
    "# Instantiate classifier\n",
    "classifier_types = ['naive Bayes']\n",
    "for clf_type in classifier_types:\n",
    "    clf = ExfiltrationClassifier(\n",
    "        clf_type, ds, verbose=False\n",
    "    )\n",
    "    # Train only on unmodified examples?\n",
    "    clf.train(include_modified=True)\n",
    "\n",
    "    # Evaluate\n",
    "    print(clf_type)\n",
    "    evaluate(clf, pp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('dns')",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "dbfbe1e77da02e8bb71d8b3a60011b526533b817325df80c6a58d482f7863cf9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}