splunk-security_content/notebooks/pretrained_dga_model_dsdl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Learning model to detect DGA domains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook uses a pre-trained deep learning model to predict whether a domain is DGA generated or not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What are DGA domains ?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adversaries may make use of Domain Generation Algorithms (DGAs) to dynamically identify a destination domain for command and control traffic rather than relying on a list of static IP addresses or domains. This has the advantage of making it much harder for defenders to block, track, or take over the command and control channel, as there potentially could be thousands of domains that malware can check for instructions.\n",
    "\n",
    "DGAs can take the form of apparently random or \"gibberish\" strings (ex: istgmxdejdnxuyla.ru) when they construct domain names by generating each letter. Alternatively, some DGAs employ whole words as the unit by concatenating words together instead of letters (ex: cityjulydish.net). Many DGAs are time-based, generating a different domain for each time period (hourly, daily, monthly, etc). Others incorporate a seed value as well to make predicting future domains more difficult for defenders https://attack.mitre.org/techniques/T1568/002/"
   ]
  },
  {
   "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": null,
   "metadata": {
    "deletable": false,
    "name": "mltkc_import"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import Input \n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.layers import Embedding\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "from tensorflow.keras.layers import Embedding,LSTM,Dropout,Dense,Activation\n",
    "from tensorflow.keras.activations import sigmoid\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "import pickle\n",
    "from sklearn.utils import shuffle\n",
    "from tensorflow.keras.preprocessing import sequence\n",
    "import datetime\n",
    "from collections import Counter\n",
    "MODEL_DIRECTORY = \"/srv/app/model/data/\"\n",
    "MAX_VOCAB = 40\n",
    "pd.options.mode.chained_assignment = None \n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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__)\n",
    "print(\"tensorflow version: \" + tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1 - get a data sample from Splunk\n",
    "This notebook does not train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "df, param = stage(\"pretrained_dga_model_dsdl\")\n",
    "print(df.describe())\n",
    "print(param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 2 - create and initialize a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "name": "mltkc_init"
   },
   "outputs": [],
   "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",
    "def init(df,param):\n",
    "    model = tf.keras.models.load_model(MODEL_DIRECTORY + \"pretrained_dga_model_dsdl\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "print(fit(model,df,param))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 4 - apply the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "name": "mltkc_apply"
   },
   "outputs": [],
   "source": [
    "# apply your model\n",
    "# returns the calculated results\n",
    "tokenizer = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/tokenizer\", 'rb'))\n",
    "vectorizer_words = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/vectorizer_words\", 'rb'))\n",
    "vectorizer_non_dga_domains = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/vectorizer_non_dga_domains\", 'rb'))\n",
    "weight_words = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/weight_words\", 'rb'))\n",
    "weight_non_dga_grams = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/weight_non_dga_grams\", 'rb'))\n",
    "domains = pickle.load(open(MODEL_DIRECTORY+\"pretrained_dga_model_dsdl/domains\", 'rb'))\n",
    "model = tf.keras.models.load_model(MODEL_DIRECTORY + \"pretrained_dga_model_dsdl\")\n",
    "\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",
    "alexa_domains = domains['domain']\n",
    "\n",
    "    \n",
    "def is_in_alexa1m(domain,domains):\n",
    "    test_in_alexa_domains = set(alexa_domains) & domains\n",
    "    return ((domain in test_in_alexa_domains))\n",
    "\n",
    "    \n",
    "    \n",
    "def add_features(df):\n",
    "    print (\"1. Done adding ngram features\")\n",
    "    X_1= weight_words * vectorizer_words.transform(df['domain']).T \n",
    "    X_2= weight_non_dga_grams * vectorizer_non_dga_domains.transform(df['domain']).T\n",
    "    X_3 =   df['domain'].map(lambda x: entropy(x)) \n",
    "    print (\"2. Done adding entropy\")\n",
    "    X_4 = df['domain'].map(lambda x: len(x))\n",
    "    print (\"3. Done adding length of domain\")\n",
    "    domains = set(df['domain'])\n",
    "    X_5 = df['domain'].map(lambda x: is_in_alexa1m(x,domains))\n",
    "    print (\"4. Done adding domain present in alexa domains\")\n",
    "    X_5 = X_5.astype(int)\n",
    "    input2 = np.c_[X_1,X_2,X_3,X_4,X_5] #\n",
    "    print (\"appending done\")\n",
    "    return input2 \n",
    "\n",
    "def prep_text(texts):\n",
    "    text_sequences = tokenizer.texts_to_sequences(texts)\n",
    "    return sequence.pad_sequences(text_sequences, maxlen=MAX_VOCAB)\n",
    "\n",
    "def apply(model,df,param):\n",
    "    input1 = prep_text(df['domain'])\n",
    "    input2 = add_features(df)\n",
    "    \n",
    "    yhat = model.predict([input1,input2])\n",
    "    y_ = (yhat > 0.5).astype(\"int32\")\n",
    "    output = pd.DataFrame()\n",
    "    #output['domain']  = df['domain']\n",
    "    #output['true_dga']  = df['is_dga'] \n",
    "    output['pred_dga']=y_.ravel()\n",
    "    output['pred_dga_proba']=yhat.ravel()\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# THIS CELL IS NOT EXPORTED - free notebook cell for testing or development purposes\n",
    "print(apply(model,df,param))"
   ]
  },
  {
   "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": null,
   "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 = tf.keras.models.load_model(MODEL_DIRECTORY + \"pretrained_dga_model_dsdl\")\n",
    "    return model"
   ]
  },
  {
   "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": [
    "## 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",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}