BSidesLV -- your model isn't that special -- (2) End-to-End.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=1 # limit GPU usage, if any to this GPU\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=1 # limit GPU usage, if any to this GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from classifier import common\n",
    "import os\n",
    "labels = common.fetch_samples()\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "np.random.seed(123)\n",
    "y_train, y_test, sha256_train, sha256_test = train_test_split(\n",
    "    list(labels.values()), list(labels.keys()), test_size=1000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# End-to-end deep learning for malware\n",
    "\n",
    "So, let's move to \"real\" end-to-end deep learning, because deep learning does everything better, right?  For this, there's a few things to note.\n",
    "1. Images have a notion of pixel intensity which has a natural ordering: black < gray < white.  But binaries are made of bytes, some of which represents instructions, some ASCII or unicode text (depending on placement and context in the file), and the byte value has no real logical ordering.  In our model, we'll let the model choose a mapping from byte value to \"color\", where you can specify how many dimensions make up each \"color\".  The layer that does this is called an embedding layer.\n",
    "2. Even hefty GPUs may have a hard time holding lots of embedded binary files in contiguous memory.  For this reason, we'll set a cap on the maximum number of bytes we'll read in (hint: the median size of malware on VirusTotal is close to 1MB), and furthermore, break the file into chunks that the memory manager can conveniently place into contiguous memory on the GPU.  (Even though you may have 4GB of memory on your GPU, you may not be able to hold a 1MB file with its embedded/\"colored\" bytes in contiguous memory.)\n",
    "3. Speaking of memory usage, we can limit the memory usage by the batch size.\n",
    "\n",
    "But, **think of it**!  End-to-end deep learing for static malware detection. **No PE parsing required!** **No feature engineering required!** **No work required!**  Right?\n",
    "\n",
    "You can find code that defines the end-to-end model architecture at [classifier/endtoend.py](classifier/endtoend.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for this demo, will slurp in only the first 256K (2**18) bytes of the file\n",
    "# for a nice GPU like a Titan X, you should be able to squeeze in > 2MB ...\n",
    "# ...but warning! this makes training more difficult...larger haystack to find needles\n",
    "max_file_length = int(2**18) # powers of 2 FTW\n",
    "file_chunks = 8  # break file into this many chunks\n",
    "file_chunk_size = max_file_length // file_chunks\n",
    "\n",
    "batch_size = 8 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100\n",
      "12375/12375 [==============================] - 3953s - loss: 0.7266 - acc: 0.7699 - val_loss: 0.5830 - val_acc: 0.8290\n",
      "Epoch 2/100\n",
      "12375/12375 [==============================] - 3875s - loss: 0.6587 - acc: 0.8010 - val_loss: 0.5958 - val_acc: 0.8500\n",
      "Epoch 3/100\n",
      "12375/12375 [==============================] - 3812s - loss: 0.6289 - acc: 0.8124 - val_loss: 0.6909 - val_acc: 0.7580\n",
      "Epoch 4/100\n",
      "12375/12375 [==============================] - 3784s - loss: 0.6077 - acc: 0.8188 - val_loss: 0.5648 - val_acc: 0.8560\n",
      "Epoch 5/100\n",
      "12375/12375 [==============================] - 3787s - loss: 0.5892 - acc: 0.8265 - val_loss: 0.5707 - val_acc: 0.7880\n",
      "Epoch 6/100\n",
      "12375/12375 [==============================] - 3787s - loss: 0.5719 - acc: 0.8288 - val_loss: 0.5401 - val_acc: 0.8450\n",
      "Epoch 7/100\n",
      "12375/12375 [==============================] - 3801s - loss: 0.5582 - acc: 0.8316 - val_loss: 0.5857 - val_acc: 0.8120\n",
      "Epoch 8/100\n",
      "12375/12375 [==============================] - 3908s - loss: 0.5420 - acc: 0.8368 - val_loss: 0.4487 - val_acc: 0.8690\n",
      "Epoch 9/100\n",
      "12375/12375 [==============================] - 3912s - loss: 0.5269 - acc: 0.8395 - val_loss: 0.4956 - val_acc: 0.8430\n",
      "Epoch 10/100\n",
      "12375/12375 [==============================] - 3806s - loss: 0.5169 - acc: 0.8392 - val_loss: 0.4832 - val_acc: 0.8840\n",
      "Epoch 11/100\n",
      "12375/12375 [==============================] - 3786s - loss: 0.5047 - acc: 0.8443 - val_loss: 0.6340 - val_acc: 0.6940\n",
      "Epoch 12/100\n",
      "12375/12375 [==============================] - 3789s - loss: 0.4930 - acc: 0.8466 - val_loss: 0.4439 - val_acc: 0.8690\n",
      "Epoch 13/100\n",
      "12375/12375 [==============================] - 3787s - loss: 0.4837 - acc: 0.8476 - val_loss: 0.4131 - val_acc: 0.8760\n",
      "Epoch 14/100\n",
      "12375/12375 [==============================] - 3791s - loss: 0.4722 - acc: 0.8503 - val_loss: 0.7371 - val_acc: 0.7790\n",
      "Epoch 15/100\n",
      "12375/12375 [==============================] - 3785s - loss: 0.4612 - acc: 0.8523 - val_loss: 0.5839 - val_acc: 0.8230\n",
      "Epoch 16/100\n",
      " 8288/12375 [===================>..........] - ETA: 1241s - loss: 0.4561 - acc: 0.8543"
     ]
    }
   ],
   "source": [
    "# That this is very long running cell, and we're going\n",
    "# it may appear that the output is truncated before training completes\n",
    "\n",
    "# let's train this puppy\n",
    "from classifier import endtoend\n",
    "import math\n",
    "from keras.callbacks import LearningRateScheduler, EarlyStopping, ModelCheckpoint, ReduceLROnPlateau\n",
    "\n",
    "# create_model(input_shape, byte_embedding_size=2, input_dropout=0.05, hidden_dropout=0.05, kernel_size=16, n_filters_per_layer=[64,256,1024], n_mlp_layers=2 )\n",
    "model_e2e = endtoend.create_model(input_shape=(file_chunks, file_chunk_size))\n",
    "train_generator = common.generator(list(zip(sha256_train, y_train)), batch_size, file_chunks, file_chunk_size)\n",
    "test_generator = common.generator(list(zip(sha256_test, y_test)), 1, file_chunks, file_chunk_size)\n",
    "training_history = model_e2e.fit_generator(train_generator,\n",
    "                        steps_per_epoch=math.ceil(len(sha256_train) / batch_size),\n",
    "                        epochs=100,\n",
    "                        callbacks=[\n",
    "                            EarlyStopping( patience=10 ),\n",
    "                            ModelCheckpoint( 'endtoend.h5', save_best_only=True),\n",
    "                            ReduceLROnPlateau( patience=5)],\n",
    "                        validation_data=test_generator,\n",
    "                        validation_steps=len(sha256_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the output above is truncated, because Jupyter notebook client couldn't muster the patience to wait for all the output coming from the kernel.  Got bored.  Moved along.  (shakes fist).  *Millennials!*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "** End-to-end convnet **\n",
      "ROC AUC = 0.949816612210904\n",
      "threshold=0.9198635220527649: 0.5020661157024794 TP rate @ 0.009689922480620155 FP rate\n",
      "confusion matrix @ threshold:\n",
      "[[511   5]\n",
      " [242 242]]\n",
      "accuracy @ threshold = 0.753\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.94981661221090397,\n",
       " 0.91986352,\n",
       " 0.0096899224806201549,\n",
       " 0.50206611570247939,\n",
       " array([[511,   5],\n",
       "        [242, 242]]),\n",
       " 0.753)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from keras.models import load_model\n",
    "# we'll load the \"best\" model (in this case, the penultimate model) that was saved \n",
    "# by our ModelCheckPoint callback\n",
    "model_e2e = load_model('endtoend.h5')\n",
    "# we could load the \"best\" model, but in this case, the \"best\" model is the penultimate, and not much better\n",
    "# than the model we have in hand\n",
    "y_pred = []\n",
    "for sha256, lab in zip(sha256_test, y_test):\n",
    "    y_pred.append(\n",
    "        model_e2e.predict_on_batch(\n",
    "            np.asarray([common.get_file_data(sha256, lab, max_file_length)]).reshape(\n",
    "                (-1, file_chunks, file_chunk_size))\n",
    "        )\n",
    "    )\n",
    "common.summarize_performance(np.asarray(y_pred).flatten(), y_test, \"End-to-end convnet\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Uggh, really?\n",
    "Wow, not really that good at all.  Looks like my fancy end-to-end model is having a hard time learning from these data.  \n",
    "I guess I need to make my model even more special?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}