notebooks/detect_dns_data_exfiltration_using_pretrained_model_in_dsdl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detect DNS Exfiltration using Deep Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook uses a pre-trained model to detect if a DNS request is a DNS exfiltration request. The pre-trained model is a deep learning model that takes in current DNS request along with previous 'x' number of requests between the same src and subdomain. The deep learning layers consume text input along with additional features such as entropy, length for the current request and aggregated features such as entropy_avg and size_avg in past few requests. \n",
    "\n",
    "The model outputs a probability score and the threshold for determining whether it is DNS exfiltration request or not is set at 0.5. Higher the `pred_is_dns_data_exfiltration_proba` probability it is more likely the DNS request is exfiltration.  The notebook has apply() method implemented and fit() method unimplemented since it uses pre-trained model for making predictions. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: By default every time you save this notebook the cells are exported into a python module which is then invoked by Splunk MLTK commands like <code> | fit ... | apply ... | summary </code>. Please read the Model Development Guide in the Deep Learning Toolkit app for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 0 - import libraries\n",
    "At stage 0 we define all imports necessary to run our subsequent code depending on various libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "deletable": false,
    "name": "mltkc_import"
   },
   "outputs": [],
   "source": [
    "# this definition exposes all python module imports that should be available in all subsequent commands\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader,TensorDataset,Dataset\n",
    "from torch.autograd import Variable as V\n",
    "import time\n",
    "from sklearn.metrics import classification_report, confusion_matrix,roc_curve, auc\n",
    "from sklearn.model_selection import train_test_split\n",
    "import string\n",
    "import torch.nn as nn\n",
    "import collections\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.preprocessing import Normalizer, MinMaxScaler\n",
    "from sklearn.pipeline import Pipeline\n",
    "from collections import Counter\n",
    "import torch.optim as optim\n",
    "import math\n",
    "MODEL_DIRECTORY = \"/srv/app/model/data/detect_dns_data_exfiltration_using_pretrained_model_in_dsdl/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy version: 1.22.1\n",
      "pandas version: 1.5.1\n"
     ]
    }
   ],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "print(\"numpy version: \" + np.__version__)\n",
    "print(\"pandas version: \" + pd.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1 - get a data sample from Splunk\n",
    "In Splunk run a search to pipe a dataset into your notebook environment. Note: mode=stage is used in the | fit command to do this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| makeresults count=10<br>\n",
    "| streamstats c as i<br>\n",
    "| eval s = i%3<br>\n",
    "| eval feature_{s}=0<br>\n",
    "| foreach feature_* [eval &lt;&lt;FIELD&gt;&gt;=random()/pow(2,31)]<br>\n",
    "| fit MLTKContainer mode=stage algo=barebone epochs=10 batch_size=1 s from feature_* into app:barebone_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you run this search your data set sample is available as a csv inside the container to develop your model. The name is taken from the into keyword (\"barebone_model\" in the example above) or set to \"default\" if no into keyword is present. This step is intended to work with a subset of your data to create your custom model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "deletable": false,
    "name": "mltkc_stage"
   },
   "outputs": [],
   "source": [
    "# this cell is not executed from MLTK and should only be used for staging data into the notebook environment\n",
    "def stage(name):\n",
    "    with open(\"data/\"+name+\".csv\", 'r') as f:\n",
    "        df = pd.read_csv(f)\n",
    "    with open(\"data/\"+name+\".json\", 'r') as f:\n",
    "        param = json.load(f)\n",
    "    return df, param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       is_exfiltration       rank\n",
      "count        10.000000  10.000000\n",
      "mean          4.800000   4.800000\n",
      "std           3.084009   3.084009\n",
      "min           1.000000   1.000000\n",
      "25%           2.000000   2.000000\n",
      "50%           4.500000   4.500000\n",
      "75%           7.750000   7.750000\n",
      "max           9.000000   9.000000\n",
      "{'options': {'params': {'mode': 'stage', 'algo': 'dns_exfil5'}, 'args': ['is_exfiltration', 'src', 'query', 'rank'], 'target_variable': ['is_exfiltration'], 'feature_variables': ['src', 'query', 'rank'], 'model_name': 'dns_exfil5', 'algo_name': 'MLTKContainer', 'mlspl_limits': {'handle_new_cat': 'default', 'max_distinct_cat_values': '100', 'max_distinct_cat_values_for_classifiers': '100', 'max_distinct_cat_values_for_scoring': '100', 'max_fit_time': '600', 'max_inputs': '100000', 'max_memory_usage_mb': '4000', 'max_model_size_mb': '30', 'max_score_time': '600', 'use_sampling': 'true'}, 'kfold_cv': None}, 'feature_variables': ['src', 'query', 'rank'], 'target_variables': ['is_exfiltration']}\n"
     ]
    }
   ],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "df, param = stage(\"dns_exfil5\")\n",
    "print(df.describe())\n",
    "print(param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 2 - create and initialize a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "deletable": false,
    "name": "mltkc_init"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "94\n"
     ]
    }
   ],
   "source": [
    "# initialize your model\n",
    "# available inputs: data and parameters\n",
    "# returns the model object which will be used as a reference to call fit, apply and summary subsequently\n",
    "device='cpu'\n",
    "class DNSExfiltration(nn.Module):\n",
    "    def __init__(self,input_size):\n",
    "        super().__init__()\n",
    "        self.layer_1 = nn.Linear(input_size, 128) \n",
    "        self.layer_2 = nn.Linear(128, 128)\n",
    "        self.layer_out = nn.Linear(128, 1) \n",
    "        \n",
    "        \n",
    "        self.relu = nn.ReLU()\n",
    "        self.sigmoid =  nn.Sigmoid()\n",
    "        self.dropout = nn.Dropout(p=0.5)\n",
    "        \n",
    "\n",
    "    def forward(self, input):\n",
    "        x = self.relu(self.layer_1(input))\n",
    "        x = self.dropout(x)\n",
    "        x = self.relu(self.layer_2(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.sigmoid(self.layer_out(x))\n",
    "        return x\n",
    "    \n",
    "\n",
    "keys = list(string.printable.strip())\n",
    "print (len(list(string.printable.strip())))\n",
    "    \n",
    "def init(df,param):\n",
    "    model = DNSExfiltration(98)\n",
    "    model.load_state_dict(torch.load(MODEL_DIRECTORY+'detect_dns_data_exfiltration_using_pretrained_model_in_dsdl.pt',map_location=torch.device('cpu')))\n",
    "    model = model.to('cpu')\n",
    "    model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DNSExfiltration(\n",
      "  (layer_1): Linear(in_features=99, out_features=256, bias=True)\n",
      "  (layer_2): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (layer_out): Linear(in_features=256, out_features=1, bias=True)\n",
      "  (relu): ReLU()\n",
      "  (sigmoid): Sigmoid()\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "print(init(df,param))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 3 - fit the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "deletable": false,
    "name": "mltkc_fit"
   },
   "outputs": [],
   "source": [
    "# train your model\n",
    "# returns a fit info json object and may modify the model object\n",
    "def fit(model,df,param):\n",
    "    # model.fit()\n",
    "    info = {\"message\": \"model trained\"}\n",
    "    return info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'message': 'model trained'}\n"
     ]
    }
   ],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "print(fit(None,None,None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 4 - apply the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "deletable": false,
    "name": "mltkc_apply"
   },
   "outputs": [],
   "source": [
    "# apply your model\n",
    "# returns the calculated results\n",
    "text_rows = []\n",
    "size_avg = []\n",
    "entropy_avg = []\n",
    "\n",
    "# this method accepts a dataloader and makes predictions in batches\n",
    "def predict(dataloader,model):\n",
    "        predict_label_proba = []\n",
    "        predict_label = []\n",
    "        for batch in (dataloader):\n",
    "            \n",
    "            #convert to 1d tensor\n",
    "            predictions = model(batch.to('cpu'))\n",
    "            output  = (predictions >= 0.5).int() \n",
    "            predict_label_proba.extend(predictions)\n",
    "            predict_label.extend(output)\n",
    "        predict_label = [x.cpu().detach().numpy().item() for x in predict_label]\n",
    "        predict_label_proba = [x.cpu().detach().numpy().item() for x in predict_label_proba]\n",
    "        return predict_label_proba,predict_label\n",
    "    \n",
    "# this method accepts a DNS request and converts into indexes based on printable characters\n",
    "def index_chars(x):\n",
    "    request_chars = {}\n",
    "    for i in range(len(x)):\n",
    "        request_chars[keys.index(x[i])] = request_chars.get(keys.index(x[i]), 0) + 1\n",
    "    text_rows.append(request_chars)\n",
    "\n",
    "#  calculates entropy of a domain\n",
    "def entropy(domain):\n",
    "    p, lns = Counter(domain), float(len(domain))\n",
    "    return -sum(count / lns * math.log(count / lns, 2) for count in p.values())\n",
    "\n",
    "\n",
    "# removes the subdomain/domain from the request\n",
    "def replace_tld(x):\n",
    "    if x is not None:\n",
    "        return str(x).rsplit('.', 2)[0]\n",
    "    else:\n",
    "        return x\n",
    "    \n",
    "# get the subdomain/domain from the request\n",
    "def get_tld(x):\n",
    "    without_tld = str(x).rsplit('.', 2)[0]\n",
    "    return str(x).replace(without_tld,'').lstrip(\".\")\n",
    "\n",
    "# compute aggregated features for the same src and subdomain/domain on a window of 10 events\n",
    "def get_aggregated_features(row,df):\n",
    "    src = row['src']\n",
    "    tld = row['tld']\n",
    "    prev_events = df[(df['src']==src) & (df['tld']==tld)]\n",
    "    \n",
    "    size_avg.append(prev_events['len'].mean())\n",
    "    entropy_avg.append(prev_events['entropy'].mean())\n",
    "\n",
    "# prepare input df by calculating features\n",
    "def prepare_input_df(df):\n",
    "    keys = list(string.printable.strip())\n",
    "\n",
    "    \n",
    "    \n",
    "    df['query'].apply(lambda x: index_chars(x))\n",
    "    text = pd.DataFrame(text_rows, columns=list(range(0, 94)))\n",
    "    text.reset_index(drop=True, inplace=True)\n",
    "    df.reset_index(drop=True, inplace=True)\n",
    "    text.fillna(0, inplace=True)\n",
    "    df = pd.concat([text, df], axis=1)\n",
    "\n",
    "\n",
    "    # request without tld\n",
    "    df['request_without_domain'] = df['query'].apply(lambda row: replace_tld(row))\n",
    "\n",
    "    # request without tld\n",
    "    df['tld'] = df['query'].apply(lambda row: get_tld(row))\n",
    "    \n",
    "    # length of domain\n",
    "    df['len'] = df['request_without_domain'].apply(len)\n",
    "    \n",
    "    # entropy\n",
    "    df['entropy'] = df['request_without_domain'].apply(lambda x: entropy(x))\n",
    "    \n",
    "    # take most-recent request\n",
    "    recent_df = df.loc[df['rank'] == 1]\n",
    "\n",
    "    # calculate feature by aggregating events\n",
    "\n",
    "    recent_df.apply(lambda x: get_aggregated_features(x,df),axis=1)\n",
    "    recent_df['size_avg'] = size_avg\n",
    "    recent_df['entropy_avg'] = entropy_avg\n",
    "    return recent_df\n",
    "\n",
    "    \n",
    "# apply model on processed dataframe to predict exfiltration\n",
    "def apply(model,df,param):\n",
    "    df.drop(['_time'], axis=1,inplace=True, errors='ignore')\n",
    "    recent_df = prepare_input_df(df)\n",
    "    input_df = recent_df.drop(['src' ,'query','rank','request_without_domain','tld'], axis=1)\n",
    "    recent_df.drop(['request_without_domain','tld','len','entropy','size_avg','entropy_avg'], axis=1, inplace=True)\n",
    "    recent_df.drop(range(0, 94),axis=1,inplace=True)\n",
    "    input_tensor = torch.FloatTensor(input_df.values)\n",
    "    dataloader = DataLoader(input_tensor, shuffle=True, batch_size=256)\n",
    "    predict_is_exfiltration_proba, predict_is_exfiltration = predict(dataloader,model)\n",
    "    recent_df['pred_is_dns_data_exfiltration_proba'] = predict_is_exfiltration_proba\n",
    "    recent_df['pred_is_dns_data_exfiltration'] = predict_is_exfiltration\n",
    "    print(recent_df.columns)\n",
    "    print(df.columns)\n",
    "    text_rows.clear()\n",
    "    size_avg.clear()\n",
    "    entropy_avg.clear()\n",
    "    output = pd.merge(recent_df,df,on=['src','query','rank'],how='right')\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 5 - save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "name": "mltkc_save"
   },
   "outputs": [],
   "source": [
    "# save model to name in expected convention \"<algo_name>_<model_name>\"\n",
    "def save(model,name):\n",
    "    with open(MODEL_DIRECTORY + name + \".json\", 'w') as file:\n",
    "        json.dump(model, file)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 6 - load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "deletable": false,
    "name": "mltkc_load"
   },
   "outputs": [],
   "source": [
    "# load model from name in expected convention \"<algo_name>_<model_name>\"\n",
    "def load(name):\n",
    "    model = DNSExfiltration(98)\n",
    "    model.load_state_dict(torch.load(MODEL_DIRECTORY+'detect_dns_data_exfiltration_using_pretrained_model_in_dsdl.pt',map_location=torch.device('cpu')))\n",
    "    model = model.to('cpu')\n",
    "    model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DNSExfiltration(\n",
       "  (layer_1): Linear(in_features=99, out_features=256, bias=True)\n",
       "  (layer_2): Linear(in_features=256, out_features=256, bias=True)\n",
       "  (layer_out): Linear(in_features=256, out_features=1, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (sigmoid): Sigmoid()\n",
       "  (dropout): Dropout(p=0.5, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "load(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 7 - provide a summary of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "name": "mltkc_summary"
   },
   "outputs": [],
   "source": [
    "# return a model summary\n",
    "def summary(model=None):\n",
    "    returns = {\"version\": {\"numpy\": np.__version__, \"pandas\": pd.__version__} }\n",
    "    return returns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After implementing your fit, apply, save and load you can train your model:<br>\n",
    "| makeresults count=10<br>\n",
    "| streamstats c as i<br>\n",
    "| eval s = i%3<br>\n",
    "| eval feature_{s}=0<br>\n",
    "| foreach feature_* [eval &lt;&lt;FIELD&gt;&gt;=random()/pow(2,31)]<br>\n",
    "| fit MLTKContainer algo=barebone s from feature_* into app:barebone_model<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or apply your model:<br>\n",
    "| makeresults count=10<br>\n",
    "| streamstats c as i<br>\n",
    "| eval s = i%3<br>\n",
    "| eval feature_{s}=0<br>\n",
    "| foreach feature_* [eval &lt;&lt;FIELD&gt;&gt;=random()/pow(2,31)]<br>\n",
    "| apply barebone_model as the_meaning_of_life"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End of Stages\n",
    "All subsequent cells are not tagged and can be used for further freeform code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}