malware-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Machine Learning approach for Malware Detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing all the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas\n",
    "import numpy\n",
    "import pickle\n",
    "import pefile\n",
    "import sklearn.ensemble as ek\n",
    "from sklearn import cross_validation, tree, linear_model\n",
    "from sklearn.feature_selection import SelectFromModel\n",
    "from sklearn.externals import joblib\n",
    "from sklearn.naive_bayes import GaussianNB\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn import preprocessing\n",
    "from sklearn import svm\n",
    "from sklearn.linear_model import LinearRegression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading the initial dataset delimited by | "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dataset = pandas.read_csv('/home/surajr/Downloads/data.csv',sep='|', low_memory=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>md5</th>\n",
       "      <th>Machine</th>\n",
       "      <th>SizeOfOptionalHeader</th>\n",
       "      <th>Characteristics</th>\n",
       "      <th>MajorLinkerVersion</th>\n",
       "      <th>MinorLinkerVersion</th>\n",
       "      <th>SizeOfCode</th>\n",
       "      <th>SizeOfInitializedData</th>\n",
       "      <th>SizeOfUninitializedData</th>\n",
       "      <th>...</th>\n",
       "      <th>ResourcesNb</th>\n",
       "      <th>ResourcesMeanEntropy</th>\n",
       "      <th>ResourcesMinEntropy</th>\n",
       "      <th>ResourcesMaxEntropy</th>\n",
       "      <th>ResourcesMeanSize</th>\n",
       "      <th>ResourcesMinSize</th>\n",
       "      <th>ResourcesMaxSize</th>\n",
       "      <th>LoadConfigurationSize</th>\n",
       "      <th>VersionInformationSize</th>\n",
       "      <th>legitimate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>memtest.exe</td>\n",
       "      <td>631ea355665f28d4707448e442fbf5b8</td>\n",
       "      <td>332</td>\n",
       "      <td>224</td>\n",
       "      <td>258</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>361984</td>\n",
       "      <td>115712</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>4</td>\n",
       "      <td>3.262823</td>\n",
       "      <td>2.568844</td>\n",
       "      <td>3.537939</td>\n",
       "      <td>8797.000000</td>\n",
       "      <td>216</td>\n",
       "      <td>18032</td>\n",
       "      <td>0</td>\n",
       "      <td>16</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ose.exe</td>\n",
       "      <td>9d10f99a6712e28f8acd5641e3a7ea6b</td>\n",
       "      <td>332</td>\n",
       "      <td>224</td>\n",
       "      <td>3330</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>130560</td>\n",
       "      <td>19968</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>4.250461</td>\n",
       "      <td>3.420744</td>\n",
       "      <td>5.080177</td>\n",
       "      <td>837.000000</td>\n",
       "      <td>518</td>\n",
       "      <td>1156</td>\n",
       "      <td>72</td>\n",
       "      <td>18</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>setup.exe</td>\n",
       "      <td>4d92f518527353c0db88a70fddcfd390</td>\n",
       "      <td>332</td>\n",
       "      <td>224</td>\n",
       "      <td>3330</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>517120</td>\n",
       "      <td>621568</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>11</td>\n",
       "      <td>4.426324</td>\n",
       "      <td>2.846449</td>\n",
       "      <td>5.271813</td>\n",
       "      <td>31102.272727</td>\n",
       "      <td>104</td>\n",
       "      <td>270376</td>\n",
       "      <td>72</td>\n",
       "      <td>18</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>DW20.EXE</td>\n",
       "      <td>a41e524f8d45f0074fd07805ff0c9b12</td>\n",
       "      <td>332</td>\n",
       "      <td>224</td>\n",
       "      <td>258</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>585728</td>\n",
       "      <td>369152</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>10</td>\n",
       "      <td>4.364291</td>\n",
       "      <td>2.669314</td>\n",
       "      <td>6.400720</td>\n",
       "      <td>1457.000000</td>\n",
       "      <td>90</td>\n",
       "      <td>4264</td>\n",
       "      <td>72</td>\n",
       "      <td>18</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>dwtrig20.exe</td>\n",
       "      <td>c87e561258f2f8650cef999bf643a731</td>\n",
       "      <td>332</td>\n",
       "      <td>224</td>\n",
       "      <td>258</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>294912</td>\n",
       "      <td>247296</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>4.306100</td>\n",
       "      <td>3.421598</td>\n",
       "      <td>5.190603</td>\n",
       "      <td>1074.500000</td>\n",
       "      <td>849</td>\n",
       "      <td>1300</td>\n",
       "      <td>72</td>\n",
       "      <td>18</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 57 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           Name                               md5  Machine  \\\n",
       "0   memtest.exe  631ea355665f28d4707448e442fbf5b8      332   \n",
       "1       ose.exe  9d10f99a6712e28f8acd5641e3a7ea6b      332   \n",
       "2     setup.exe  4d92f518527353c0db88a70fddcfd390      332   \n",
       "3      DW20.EXE  a41e524f8d45f0074fd07805ff0c9b12      332   \n",
       "4  dwtrig20.exe  c87e561258f2f8650cef999bf643a731      332   \n",
       "\n",
       "   SizeOfOptionalHeader  Characteristics  MajorLinkerVersion  \\\n",
       "0                   224              258                   9   \n",
       "1                   224             3330                   9   \n",
       "2                   224             3330                   9   \n",
       "3                   224              258                   9   \n",
       "4                   224              258                   9   \n",
       "\n",
       "   MinorLinkerVersion  SizeOfCode  SizeOfInitializedData  \\\n",
       "0                   0      361984                 115712   \n",
       "1                   0      130560                  19968   \n",
       "2                   0      517120                 621568   \n",
       "3                   0      585728                 369152   \n",
       "4                   0      294912                 247296   \n",
       "\n",
       "   SizeOfUninitializedData     ...      ResourcesNb  ResourcesMeanEntropy  \\\n",
       "0                        0     ...                4              3.262823   \n",
       "1                        0     ...                2              4.250461   \n",
       "2                        0     ...               11              4.426324   \n",
       "3                        0     ...               10              4.364291   \n",
       "4                        0     ...                2              4.306100   \n",
       "\n",
       "   ResourcesMinEntropy ResourcesMaxEntropy  ResourcesMeanSize  \\\n",
       "0             2.568844            3.537939        8797.000000   \n",
       "1             3.420744            5.080177         837.000000   \n",
       "2             2.846449            5.271813       31102.272727   \n",
       "3             2.669314            6.400720        1457.000000   \n",
       "4             3.421598            5.190603        1074.500000   \n",
       "\n",
       "   ResourcesMinSize  ResourcesMaxSize  LoadConfigurationSize  \\\n",
       "0               216             18032                      0   \n",
       "1               518              1156                     72   \n",
       "2               104            270376                     72   \n",
       "3                90              4264                     72   \n",
       "4               849              1300                     72   \n",
       "\n",
       "   VersionInformationSize  legitimate  \n",
       "0                      16           1  \n",
       "1                      18           1  \n",
       "2                      18           1  \n",
       "3                      18           1  \n",
       "4                      18           1  \n",
       "\n",
       "[5 rows x 57 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Machine</th>\n",
       "      <th>SizeOfOptionalHeader</th>\n",
       "      <th>Characteristics</th>\n",
       "      <th>MajorLinkerVersion</th>\n",
       "      <th>MinorLinkerVersion</th>\n",
       "      <th>SizeOfCode</th>\n",
       "      <th>SizeOfInitializedData</th>\n",
       "      <th>SizeOfUninitializedData</th>\n",
       "      <th>AddressOfEntryPoint</th>\n",
       "      <th>BaseOfCode</th>\n",
       "      <th>...</th>\n",
       "      <th>ResourcesNb</th>\n",
       "      <th>ResourcesMeanEntropy</th>\n",
       "      <th>ResourcesMinEntropy</th>\n",
       "      <th>ResourcesMaxEntropy</th>\n",
       "      <th>ResourcesMeanSize</th>\n",
       "      <th>ResourcesMinSize</th>\n",
       "      <th>ResourcesMaxSize</th>\n",
       "      <th>LoadConfigurationSize</th>\n",
       "      <th>VersionInformationSize</th>\n",
       "      <th>legitimate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>...</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>1.380470e+05</td>\n",
       "      <td>138047.000000</td>\n",
       "      <td>138047.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>4259.069274</td>\n",
       "      <td>225.845632</td>\n",
       "      <td>4444.145994</td>\n",
       "      <td>8.619774</td>\n",
       "      <td>3.819286</td>\n",
       "      <td>2.425956e+05</td>\n",
       "      <td>4.504867e+05</td>\n",
       "      <td>1.009525e+05</td>\n",
       "      <td>1.719561e+05</td>\n",
       "      <td>5.779845e+04</td>\n",
       "      <td>...</td>\n",
       "      <td>22.050700</td>\n",
       "      <td>4.000127</td>\n",
       "      <td>2.434541</td>\n",
       "      <td>5.521610</td>\n",
       "      <td>5.545093e+04</td>\n",
       "      <td>1.818082e+04</td>\n",
       "      <td>2.465903e+05</td>\n",
       "      <td>4.656750e+05</td>\n",
       "      <td>12.363115</td>\n",
       "      <td>0.299340</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>10880.347245</td>\n",
       "      <td>5.121399</td>\n",
       "      <td>8186.782524</td>\n",
       "      <td>4.088757</td>\n",
       "      <td>11.862675</td>\n",
       "      <td>5.754485e+06</td>\n",
       "      <td>2.101599e+07</td>\n",
       "      <td>1.635288e+07</td>\n",
       "      <td>3.430553e+06</td>\n",
       "      <td>5.527658e+06</td>\n",
       "      <td>...</td>\n",
       "      <td>136.494244</td>\n",
       "      <td>1.112981</td>\n",
       "      <td>0.815577</td>\n",
       "      <td>1.597403</td>\n",
       "      <td>7.799163e+06</td>\n",
       "      <td>6.502369e+06</td>\n",
       "      <td>2.124860e+07</td>\n",
       "      <td>2.608987e+07</td>\n",
       "      <td>6.798878</td>\n",
       "      <td>0.457971</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>332.000000</td>\n",
       "      <td>224.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>332.000000</td>\n",
       "      <td>224.000000</td>\n",
       "      <td>258.000000</td>\n",
       "      <td>8.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>3.020800e+04</td>\n",
       "      <td>2.457600e+04</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>1.272100e+04</td>\n",
       "      <td>4.096000e+03</td>\n",
       "      <td>...</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>3.458505</td>\n",
       "      <td>2.178748</td>\n",
       "      <td>4.828706</td>\n",
       "      <td>9.560000e+02</td>\n",
       "      <td>4.800000e+01</td>\n",
       "      <td>2.216000e+03</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>13.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>332.000000</td>\n",
       "      <td>224.000000</td>\n",
       "      <td>258.000000</td>\n",
       "      <td>9.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.136640e+05</td>\n",
       "      <td>2.631680e+05</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>5.288300e+04</td>\n",
       "      <td>4.096000e+03</td>\n",
       "      <td>...</td>\n",
       "      <td>6.000000</td>\n",
       "      <td>3.729824</td>\n",
       "      <td>2.458492</td>\n",
       "      <td>5.317552</td>\n",
       "      <td>2.708154e+03</td>\n",
       "      <td>4.800000e+01</td>\n",
       "      <td>9.640000e+03</td>\n",
       "      <td>7.200000e+01</td>\n",
       "      <td>15.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>332.000000</td>\n",
       "      <td>224.000000</td>\n",
       "      <td>8226.000000</td>\n",
       "      <td>10.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.203200e+05</td>\n",
       "      <td>3.850240e+05</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>6.157800e+04</td>\n",
       "      <td>4.096000e+03</td>\n",
       "      <td>...</td>\n",
       "      <td>13.000000</td>\n",
       "      <td>4.233051</td>\n",
       "      <td>2.696833</td>\n",
       "      <td>6.502239</td>\n",
       "      <td>6.558429e+03</td>\n",
       "      <td>1.320000e+02</td>\n",
       "      <td>2.378000e+04</td>\n",
       "      <td>7.200000e+01</td>\n",
       "      <td>16.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>34404.000000</td>\n",
       "      <td>352.000000</td>\n",
       "      <td>49551.000000</td>\n",
       "      <td>255.000000</td>\n",
       "      <td>255.000000</td>\n",
       "      <td>1.818587e+09</td>\n",
       "      <td>4.294966e+09</td>\n",
       "      <td>4.294941e+09</td>\n",
       "      <td>1.074484e+09</td>\n",
       "      <td>2.028711e+09</td>\n",
       "      <td>...</td>\n",
       "      <td>7694.000000</td>\n",
       "      <td>7.999723</td>\n",
       "      <td>7.999723</td>\n",
       "      <td>8.000000</td>\n",
       "      <td>2.415919e+09</td>\n",
       "      <td>2.415919e+09</td>\n",
       "      <td>4.294903e+09</td>\n",
       "      <td>4.294967e+09</td>\n",
       "      <td>26.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>8 rows × 54 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             Machine  SizeOfOptionalHeader  Characteristics  \\\n",
       "count  138047.000000         138047.000000    138047.000000   \n",
       "mean     4259.069274            225.845632      4444.145994   \n",
       "std     10880.347245              5.121399      8186.782524   \n",
       "min       332.000000            224.000000         2.000000   \n",
       "25%       332.000000            224.000000       258.000000   \n",
       "50%       332.000000            224.000000       258.000000   \n",
       "75%       332.000000            224.000000      8226.000000   \n",
       "max     34404.000000            352.000000     49551.000000   \n",
       "\n",
       "       MajorLinkerVersion  MinorLinkerVersion    SizeOfCode  \\\n",
       "count       138047.000000       138047.000000  1.380470e+05   \n",
       "mean             8.619774            3.819286  2.425956e+05   \n",
       "std              4.088757           11.862675  5.754485e+06   \n",
       "min              0.000000            0.000000  0.000000e+00   \n",
       "25%              8.000000            0.000000  3.020800e+04   \n",
       "50%              9.000000            0.000000  1.136640e+05   \n",
       "75%             10.000000            0.000000  1.203200e+05   \n",
       "max            255.000000          255.000000  1.818587e+09   \n",
       "\n",
       "       SizeOfInitializedData  SizeOfUninitializedData  AddressOfEntryPoint  \\\n",
       "count           1.380470e+05             1.380470e+05         1.380470e+05   \n",
       "mean            4.504867e+05             1.009525e+05         1.719561e+05   \n",
       "std             2.101599e+07             1.635288e+07         3.430553e+06   \n",
       "min             0.000000e+00             0.000000e+00         0.000000e+00   \n",
       "25%             2.457600e+04             0.000000e+00         1.272100e+04   \n",
       "50%             2.631680e+05             0.000000e+00         5.288300e+04   \n",
       "75%             3.850240e+05             0.000000e+00         6.157800e+04   \n",
       "max             4.294966e+09             4.294941e+09         1.074484e+09   \n",
       "\n",
       "         BaseOfCode      ...          ResourcesNb  ResourcesMeanEntropy  \\\n",
       "count  1.380470e+05      ...        138047.000000         138047.000000   \n",
       "mean   5.779845e+04      ...            22.050700              4.000127   \n",
       "std    5.527658e+06      ...           136.494244              1.112981   \n",
       "min    0.000000e+00      ...             0.000000              0.000000   \n",
       "25%    4.096000e+03      ...             5.000000              3.458505   \n",
       "50%    4.096000e+03      ...             6.000000              3.729824   \n",
       "75%    4.096000e+03      ...            13.000000              4.233051   \n",
       "max    2.028711e+09      ...          7694.000000              7.999723   \n",
       "\n",
       "       ResourcesMinEntropy  ResourcesMaxEntropy  ResourcesMeanSize  \\\n",
       "count        138047.000000        138047.000000       1.380470e+05   \n",
       "mean              2.434541             5.521610       5.545093e+04   \n",
       "std               0.815577             1.597403       7.799163e+06   \n",
       "min               0.000000             0.000000       0.000000e+00   \n",
       "25%               2.178748             4.828706       9.560000e+02   \n",
       "50%               2.458492             5.317552       2.708154e+03   \n",
       "75%               2.696833             6.502239       6.558429e+03   \n",
       "max               7.999723             8.000000       2.415919e+09   \n",
       "\n",
       "       ResourcesMinSize  ResourcesMaxSize  LoadConfigurationSize  \\\n",
       "count      1.380470e+05      1.380470e+05           1.380470e+05   \n",
       "mean       1.818082e+04      2.465903e+05           4.656750e+05   \n",
       "std        6.502369e+06      2.124860e+07           2.608987e+07   \n",
       "min        0.000000e+00      0.000000e+00           0.000000e+00   \n",
       "25%        4.800000e+01      2.216000e+03           0.000000e+00   \n",
       "50%        4.800000e+01      9.640000e+03           7.200000e+01   \n",
       "75%        1.320000e+02      2.378000e+04           7.200000e+01   \n",
       "max        2.415919e+09      4.294903e+09           4.294967e+09   \n",
       "\n",
       "       VersionInformationSize     legitimate  \n",
       "count           138047.000000  138047.000000  \n",
       "mean                12.363115       0.299340  \n",
       "std                  6.798878       0.457971  \n",
       "min                  0.000000       0.000000  \n",
       "25%                 13.000000       0.000000  \n",
       "50%                 15.000000       0.000000  \n",
       "75%                 16.000000       1.000000  \n",
       "max                 26.000000       1.000000  \n",
       "\n",
       "[8 rows x 54 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Number of malicious files vs Legitimate files in the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "legitimate\n",
       "0    96724\n",
       "1    41323\n",
       "dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.groupby(dataset['legitimate']).size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dropping columns like Name of the file, MD5 (message digest) and label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "X = dataset.drop(['Name','md5','legitimate'],axis=1).values\n",
    "y = dataset['legitimate'].values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### ExtraTreesClassifier\n",
    "ExtraTreesClassifier fits a number of randomized decision trees (a.k.a. extra-trees) on various sub-samples of the dataset and use averaging to improve the predictive accuracy and control over-fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "extratrees = ek.ExtraTreesClassifier().fit(X,y)\n",
    "model = SelectFromModel(extratrees, prefit=True)\n",
    "X_new = model.transform(X)\n",
    "nbfeatures = X_new.shape[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ExtraTreesClassifier helps in selecting the required features useful for classifying a file as either Malicious or Legitimate\n",
    "\n",
    "14 features are identified as required by ExtraTreesClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nbfeatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "######  Cross Validation\n",
    "Cross validation is applied to divide the dataset into random train and test subsets.\n",
    "test_size = 0.2 represent the proportion of the dataset to include in the test split "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = cross_validation.train_test_split(X_new, y ,test_size=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features = []\n",
    "index = numpy.argsort(extratrees.feature_importances_)[::-1][:nbfeatures]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The features identified by ExtraTreesClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. feature DllCharacteristics (0.141259)\n",
      "2. feature Characteristics (0.136174)\n",
      "3. feature Machine (0.102237)\n",
      "4. feature SectionsMaxEntropy (0.093866)\n",
      "5. feature MajorSubsystemVersion (0.076185)\n",
      "6. feature ResourcesMinEntropy (0.054568)\n",
      "7. feature ResourcesMaxEntropy (0.048843)\n",
      "8. feature ImageBase (0.047034)\n",
      "9. feature VersionInformationSize (0.046712)\n",
      "10. feature SizeOfOptionalHeader (0.041392)\n",
      "11. feature SectionsMeanEntropy (0.025279)\n",
      "12. feature Subsystem (0.022657)\n",
      "13. feature MajorOperatingSystemVersion (0.019587)\n",
      "14. feature CheckSum (0.019544)\n"
     ]
    }
   ],
   "source": [
    "for f in range(nbfeatures):\n",
    "    print(\"%d. feature %s (%f)\" % (f + 1, dataset.columns[2+index[f]], extratrees.feature_importances_[index[f]]))\n",
    "    features.append(dataset.columns[2+f])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Building the below Machine Learning model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = { \"DecisionTree\":tree.DecisionTreeClassifier(max_depth=10),\n",
    "         \"RandomForest\":ek.RandomForestClassifier(n_estimators=50),\n",
    "         \"Adaboost\":ek.AdaBoostClassifier(n_estimators=50),\n",
    "         \"GradientBoosting\":ek.GradientBoostingClassifier(n_estimators=50),\n",
    "         \"GNB\":GaussianNB(),\n",
    "         \"LinearRegression\":LinearRegression()   \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training each of the model with the X_train and testing with X_test.\n",
    "The model with best accuracy will be ranked as winner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RandomForest : 0.994386091996 \n",
      "GradientBoosting : 0.988373777617 \n",
      "GNB : 0.702897500905 \n",
      "DecisionTree : 0.990981528432 \n",
      "LinearRegression : 0.54036008649 \n",
      "Adaboost : 0.986381745744 \n"
     ]
    }
   ],
   "source": [
    "results = {}\n",
    "for algo in model:\n",
    "    clf = model[algo]\n",
    "    clf.fit(X_train,y_train)\n",
    "    score = clf.score(X_test,y_test)\n",
    "    print (\"%s : %s \" %(algo, score))\n",
    "    results[algo] = score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "winner = max(results, key=results.get)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Saving the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['classifier/classifier.pkl',\n",
       " 'classifier/classifier.pkl_01.npy',\n",
       " 'classifier/classifier.pkl_02.npy',\n",
       " 'classifier/classifier.pkl_03.npy',\n",
       " 'classifier/classifier.pkl_04.npy',\n",
       " 'classifier/classifier.pkl_05.npy',\n",
       " 'classifier/classifier.pkl_06.npy',\n",
       " 'classifier/classifier.pkl_07.npy',\n",
       " 'classifier/classifier.pkl_08.npy',\n",
       " 'classifier/classifier.pkl_09.npy',\n",
       " 'classifier/classifier.pkl_10.npy',\n",
       " 'classifier/classifier.pkl_11.npy',\n",
       " 'classifier/classifier.pkl_12.npy',\n",
       " 'classifier/classifier.pkl_13.npy',\n",
       " 'classifier/classifier.pkl_14.npy',\n",
       " 'classifier/classifier.pkl_15.npy',\n",
       " 'classifier/classifier.pkl_16.npy',\n",
       " 'classifier/classifier.pkl_17.npy',\n",
       " 'classifier/classifier.pkl_18.npy',\n",
       " 'classifier/classifier.pkl_19.npy',\n",
       " 'classifier/classifier.pkl_20.npy',\n",
       " 'classifier/classifier.pkl_21.npy',\n",
       " 'classifier/classifier.pkl_22.npy',\n",
       " 'classifier/classifier.pkl_23.npy',\n",
       " 'classifier/classifier.pkl_24.npy',\n",
       " 'classifier/classifier.pkl_25.npy',\n",
       " 'classifier/classifier.pkl_26.npy',\n",
       " 'classifier/classifier.pkl_27.npy',\n",
       " 'classifier/classifier.pkl_28.npy',\n",
       " 'classifier/classifier.pkl_29.npy',\n",
       " 'classifier/classifier.pkl_30.npy',\n",
       " 'classifier/classifier.pkl_31.npy',\n",
       " 'classifier/classifier.pkl_32.npy',\n",
       " 'classifier/classifier.pkl_33.npy',\n",
       " 'classifier/classifier.pkl_34.npy',\n",
       " 'classifier/classifier.pkl_35.npy',\n",
       " 'classifier/classifier.pkl_36.npy',\n",
       " 'classifier/classifier.pkl_37.npy',\n",
       " 'classifier/classifier.pkl_38.npy',\n",
       " 'classifier/classifier.pkl_39.npy',\n",
       " 'classifier/classifier.pkl_40.npy',\n",
       " 'classifier/classifier.pkl_41.npy',\n",
       " 'classifier/classifier.pkl_42.npy',\n",
       " 'classifier/classifier.pkl_43.npy',\n",
       " 'classifier/classifier.pkl_44.npy',\n",
       " 'classifier/classifier.pkl_45.npy',\n",
       " 'classifier/classifier.pkl_46.npy',\n",
       " 'classifier/classifier.pkl_47.npy',\n",
       " 'classifier/classifier.pkl_48.npy',\n",
       " 'classifier/classifier.pkl_49.npy',\n",
       " 'classifier/classifier.pkl_50.npy',\n",
       " 'classifier/classifier.pkl_51.npy',\n",
       " 'classifier/classifier.pkl_52.npy',\n",
       " 'classifier/classifier.pkl_53.npy',\n",
       " 'classifier/classifier.pkl_54.npy',\n",
       " 'classifier/classifier.pkl_55.npy',\n",
       " 'classifier/classifier.pkl_56.npy',\n",
       " 'classifier/classifier.pkl_57.npy',\n",
       " 'classifier/classifier.pkl_58.npy',\n",
       " 'classifier/classifier.pkl_59.npy',\n",
       " 'classifier/classifier.pkl_60.npy',\n",
       " 'classifier/classifier.pkl_61.npy',\n",
       " 'classifier/classifier.pkl_62.npy',\n",
       " 'classifier/classifier.pkl_63.npy',\n",
       " 'classifier/classifier.pkl_64.npy',\n",
       " 'classifier/classifier.pkl_65.npy',\n",
       " 'classifier/classifier.pkl_66.npy',\n",
       " 'classifier/classifier.pkl_67.npy',\n",
       " 'classifier/classifier.pkl_68.npy',\n",
       " 'classifier/classifier.pkl_69.npy',\n",
       " 'classifier/classifier.pkl_70.npy',\n",
       " 'classifier/classifier.pkl_71.npy',\n",
       " 'classifier/classifier.pkl_72.npy',\n",
       " 'classifier/classifier.pkl_73.npy',\n",
       " 'classifier/classifier.pkl_74.npy',\n",
       " 'classifier/classifier.pkl_75.npy',\n",
       " 'classifier/classifier.pkl_76.npy',\n",
       " 'classifier/classifier.pkl_77.npy',\n",
       " 'classifier/classifier.pkl_78.npy',\n",
       " 'classifier/classifier.pkl_79.npy',\n",
       " 'classifier/classifier.pkl_80.npy',\n",
       " 'classifier/classifier.pkl_81.npy',\n",
       " 'classifier/classifier.pkl_82.npy',\n",
       " 'classifier/classifier.pkl_83.npy',\n",
       " 'classifier/classifier.pkl_84.npy',\n",
       " 'classifier/classifier.pkl_85.npy',\n",
       " 'classifier/classifier.pkl_86.npy',\n",
       " 'classifier/classifier.pkl_87.npy',\n",
       " 'classifier/classifier.pkl_88.npy',\n",
       " 'classifier/classifier.pkl_89.npy',\n",
       " 'classifier/classifier.pkl_90.npy',\n",
       " 'classifier/classifier.pkl_91.npy',\n",
       " 'classifier/classifier.pkl_92.npy',\n",
       " 'classifier/classifier.pkl_93.npy',\n",
       " 'classifier/classifier.pkl_94.npy',\n",
       " 'classifier/classifier.pkl_95.npy',\n",
       " 'classifier/classifier.pkl_96.npy',\n",
       " 'classifier/classifier.pkl_97.npy',\n",
       " 'classifier/classifier.pkl_98.npy',\n",
       " 'classifier/classifier.pkl_99.npy',\n",
       " 'classifier/classifier.pkl_100.npy',\n",
       " 'classifier/classifier.pkl_101.npy',\n",
       " 'classifier/classifier.pkl_102.npy',\n",
       " 'classifier/classifier.pkl_103.npy',\n",
       " 'classifier/classifier.pkl_104.npy',\n",
       " 'classifier/classifier.pkl_105.npy',\n",
       " 'classifier/classifier.pkl_106.npy',\n",
       " 'classifier/classifier.pkl_107.npy',\n",
       " 'classifier/classifier.pkl_108.npy',\n",
       " 'classifier/classifier.pkl_109.npy',\n",
       " 'classifier/classifier.pkl_110.npy',\n",
       " 'classifier/classifier.pkl_111.npy',\n",
       " 'classifier/classifier.pkl_112.npy',\n",
       " 'classifier/classifier.pkl_113.npy',\n",
       " 'classifier/classifier.pkl_114.npy',\n",
       " 'classifier/classifier.pkl_115.npy',\n",
       " 'classifier/classifier.pkl_116.npy',\n",
       " 'classifier/classifier.pkl_117.npy',\n",
       " 'classifier/classifier.pkl_118.npy',\n",
       " 'classifier/classifier.pkl_119.npy',\n",
       " 'classifier/classifier.pkl_120.npy',\n",
       " 'classifier/classifier.pkl_121.npy',\n",
       " 'classifier/classifier.pkl_122.npy',\n",
       " 'classifier/classifier.pkl_123.npy',\n",
       " 'classifier/classifier.pkl_124.npy',\n",
       " 'classifier/classifier.pkl_125.npy',\n",
       " 'classifier/classifier.pkl_126.npy',\n",
       " 'classifier/classifier.pkl_127.npy',\n",
       " 'classifier/classifier.pkl_128.npy',\n",
       " 'classifier/classifier.pkl_129.npy',\n",
       " 'classifier/classifier.pkl_130.npy',\n",
       " 'classifier/classifier.pkl_131.npy',\n",
       " 'classifier/classifier.pkl_132.npy',\n",
       " 'classifier/classifier.pkl_133.npy',\n",
       " 'classifier/classifier.pkl_134.npy',\n",
       " 'classifier/classifier.pkl_135.npy',\n",
       " 'classifier/classifier.pkl_136.npy',\n",
       " 'classifier/classifier.pkl_137.npy',\n",
       " 'classifier/classifier.pkl_138.npy',\n",
       " 'classifier/classifier.pkl_139.npy',\n",
       " 'classifier/classifier.pkl_140.npy',\n",
       " 'classifier/classifier.pkl_141.npy',\n",
       " 'classifier/classifier.pkl_142.npy',\n",
       " 'classifier/classifier.pkl_143.npy',\n",
       " 'classifier/classifier.pkl_144.npy',\n",
       " 'classifier/classifier.pkl_145.npy',\n",
       " 'classifier/classifier.pkl_146.npy',\n",
       " 'classifier/classifier.pkl_147.npy',\n",
       " 'classifier/classifier.pkl_148.npy',\n",
       " 'classifier/classifier.pkl_149.npy',\n",
       " 'classifier/classifier.pkl_150.npy',\n",
       " 'classifier/classifier.pkl_151.npy',\n",
       " 'classifier/classifier.pkl_152.npy',\n",
       " 'classifier/classifier.pkl_153.npy',\n",
       " 'classifier/classifier.pkl_154.npy',\n",
       " 'classifier/classifier.pkl_155.npy',\n",
       " 'classifier/classifier.pkl_156.npy',\n",
       " 'classifier/classifier.pkl_157.npy',\n",
       " 'classifier/classifier.pkl_158.npy',\n",
       " 'classifier/classifier.pkl_159.npy',\n",
       " 'classifier/classifier.pkl_160.npy',\n",
       " 'classifier/classifier.pkl_161.npy',\n",
       " 'classifier/classifier.pkl_162.npy',\n",
       " 'classifier/classifier.pkl_163.npy',\n",
       " 'classifier/classifier.pkl_164.npy',\n",
       " 'classifier/classifier.pkl_165.npy',\n",
       " 'classifier/classifier.pkl_166.npy',\n",
       " 'classifier/classifier.pkl_167.npy',\n",
       " 'classifier/classifier.pkl_168.npy',\n",
       " 'classifier/classifier.pkl_169.npy',\n",
       " 'classifier/classifier.pkl_170.npy',\n",
       " 'classifier/classifier.pkl_171.npy',\n",
       " 'classifier/classifier.pkl_172.npy',\n",
       " 'classifier/classifier.pkl_173.npy',\n",
       " 'classifier/classifier.pkl_174.npy',\n",
       " 'classifier/classifier.pkl_175.npy',\n",
       " 'classifier/classifier.pkl_176.npy',\n",
       " 'classifier/classifier.pkl_177.npy',\n",
       " 'classifier/classifier.pkl_178.npy',\n",
       " 'classifier/classifier.pkl_179.npy',\n",
       " 'classifier/classifier.pkl_180.npy',\n",
       " 'classifier/classifier.pkl_181.npy',\n",
       " 'classifier/classifier.pkl_182.npy',\n",
       " 'classifier/classifier.pkl_183.npy',\n",
       " 'classifier/classifier.pkl_184.npy',\n",
       " 'classifier/classifier.pkl_185.npy',\n",
       " 'classifier/classifier.pkl_186.npy',\n",
       " 'classifier/classifier.pkl_187.npy',\n",
       " 'classifier/classifier.pkl_188.npy',\n",
       " 'classifier/classifier.pkl_189.npy',\n",
       " 'classifier/classifier.pkl_190.npy',\n",
       " 'classifier/classifier.pkl_191.npy',\n",
       " 'classifier/classifier.pkl_192.npy',\n",
       " 'classifier/classifier.pkl_193.npy',\n",
       " 'classifier/classifier.pkl_194.npy',\n",
       " 'classifier/classifier.pkl_195.npy',\n",
       " 'classifier/classifier.pkl_196.npy',\n",
       " 'classifier/classifier.pkl_197.npy',\n",
       " 'classifier/classifier.pkl_198.npy',\n",
       " 'classifier/classifier.pkl_199.npy',\n",
       " 'classifier/classifier.pkl_200.npy',\n",
       " 'classifier/classifier.pkl_201.npy']"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joblib.dump(model[winner],'classifier/classifier.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "open('classifier/features.pkl', 'w').write(pickle.dumps(features))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculating the False positive and negative on the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False positive rate : 0.099251 %\n",
      "False negative rate : 0.147618 %\n"
     ]
    }
   ],
   "source": [
    "clf = model[winner]\n",
    "res = clf.predict(X_new)\n",
    "mt = confusion_matrix(y, res)\n",
    "print(\"False positive rate : %f %%\" % ((mt[0][1] / float(sum(mt[0])))*100))\n",
    "print('False negative rate : %f %%' % ( (mt[1][0] / float(sum(mt[1]))*100)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Load classifier\n",
    "clf = joblib.load('classifier/classifier.pkl')\n",
    "#load features\n",
    "features = pickle.loads(open(os.path.join('classifier/features.pkl'),'r').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Testing with unseen file\n",
    "Given any unseen test file, it's required to extract the characteristics of the given file.  \n",
    "\n",
    "In order to test the model on an unseen file, it's required to extract the characteristics of the given file. Python's pefile.PE library is used to construct and build the feature vector and a ML model is used to predict the class for the given file based on the already trained model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# %load malware_test.py\n",
    "\"\"\"\n",
    "this file extracts the required information of a given file using the library PE \n",
    "\n",
    "\"\"\"\n",
    "\n",
    "import pefile\n",
    "import os\n",
    "import array\n",
    "import math\n",
    "import pickle\n",
    "from sklearn.externals import joblib\n",
    "import sys\n",
    "import argparse\n",
    "\n",
    "\n",
    "\n",
    "def get_entropy(data):\n",
    "    if len(data) == 0:\n",
    "\treturn 0.0\n",
    "    occurences = array.array('L', [0]*256)\n",
    "    for x in data:\n",
    "  \toccurences[x if isinstance(x, int) else ord(x)] += 1\n",
    "\n",
    "    entropy = 0\n",
    "    for x in occurences:\n",
    "\tif x:\n",
    "\t    p_x = float(x) / len(data)\n",
    "\t    entropy -= p_x*math.log(p_x, 2)\n",
    "\n",
    "    return entropy\n",
    "\n",
    "\n",
    "def get_resources(pe):\n",
    "    \"\"\"Extract resources :\n",
    "    [entropy, size]\"\"\"\n",
    "    resources = []\n",
    "    if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):\n",
    "\ttry:\n",
    "            for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:\n",
    "                if hasattr(resource_type, 'directory'):\n",
    "                    for resource_id in resource_type.directory.entries:\n",
    "                        if hasattr(resource_id, 'directory'):\n",
    "                            for resource_lang in resource_id.directory.entries:\n",
    "                                data = pe.get_data(resource_lang.data.struct.OffsetToData, resource_lang.data.struct.Size)\n",
    "                                size = resource_lang.data.struct.Size\n",
    "                                entropy = get_entropy(data)\n",
    "\n",
    "                                resources.append([entropy, size])\n",
    "        except Exception as e:\n",
    "            return resources\n",
    "    return resources\n",
    "\n",
    "def get_version_info(pe):\n",
    "    \"\"\"Return version infos\"\"\"\n",
    "    res = {}\n",
    "    for fileinfo in pe.FileInfo:\n",
    "        if fileinfo.Key == 'StringFileInfo':\n",
    "            for st in fileinfo.StringTable:\n",
    "                for entry in st.entries.items():\n",
    "                    res[entry[0]] = entry[1]\n",
    "        if fileinfo.Key == 'VarFileInfo':\n",
    "            for var in fileinfo.Var:\n",
    "                res[var.entry.items()[0][0]] = var.entry.items()[0][1]\n",
    "    if hasattr(pe, 'VS_FIXEDFILEINFO'):\n",
    "          res['flags'] = pe.VS_FIXEDFILEINFO.FileFlags\n",
    "          res['os'] = pe.VS_FIXEDFILEINFO.FileOS\n",
    "          res['type'] = pe.VS_FIXEDFILEINFO.FileType\n",
    "          res['file_version'] = pe.VS_FIXEDFILEINFO.FileVersionLS\n",
    "          res['product_version'] = pe.VS_FIXEDFILEINFO.ProductVersionLS\n",
    "          res['signature'] = pe.VS_FIXEDFILEINFO.Signature\n",
    "          res['struct_version'] = pe.VS_FIXEDFILEINFO.StrucVersion\n",
    "    return res\n",
    "\n",
    "#extract the info for a given file\n",
    "def extract_infos(fpath):\n",
    "    res = {}\n",
    "    pe = pefile.PE(fpath)\n",
    "    res['Machine'] = pe.FILE_HEADER.Machine\n",
    "    res['SizeOfOptionalHeader'] = pe.FILE_HEADER.SizeOfOptionalHeader\n",
    "    res['Characteristics'] = pe.FILE_HEADER.Characteristics\n",
    "    res['MajorLinkerVersion'] = pe.OPTIONAL_HEADER.MajorLinkerVersion\n",
    "    res['MinorLinkerVersion'] = pe.OPTIONAL_HEADER.MinorLinkerVersion\n",
    "    res['SizeOfCode'] = pe.OPTIONAL_HEADER.SizeOfCode\n",
    "    res['SizeOfInitializedData'] = pe.OPTIONAL_HEADER.SizeOfInitializedData\n",
    "    res['SizeOfUninitializedData'] = pe.OPTIONAL_HEADER.SizeOfUninitializedData\n",
    "    res['AddressOfEntryPoint'] = pe.OPTIONAL_HEADER.AddressOfEntryPoint\n",
    "    res['BaseOfCode'] = pe.OPTIONAL_HEADER.BaseOfCode\n",
    "    try:\n",
    "        res['BaseOfData'] = pe.OPTIONAL_HEADER.BaseOfData\n",
    "    except AttributeError:\n",
    "        res['BaseOfData'] = 0\n",
    "    res['ImageBase'] = pe.OPTIONAL_HEADER.ImageBase\n",
    "    res['SectionAlignment'] = pe.OPTIONAL_HEADER.SectionAlignment\n",
    "    res['FileAlignment'] = pe.OPTIONAL_HEADER.FileAlignment\n",
    "    res['MajorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MajorOperatingSystemVersion\n",
    "    res['MinorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MinorOperatingSystemVersion\n",
    "    res['MajorImageVersion'] = pe.OPTIONAL_HEADER.MajorImageVersion\n",
    "    res['MinorImageVersion'] = pe.OPTIONAL_HEADER.MinorImageVersion\n",
    "    res['MajorSubsystemVersion'] = pe.OPTIONAL_HEADER.MajorSubsystemVersion\n",
    "    res['MinorSubsystemVersion'] = pe.OPTIONAL_HEADER.MinorSubsystemVersion\n",
    "    res['SizeOfImage'] = pe.OPTIONAL_HEADER.SizeOfImage\n",
    "    res['SizeOfHeaders'] = pe.OPTIONAL_HEADER.SizeOfHeaders\n",
    "    res['CheckSum'] = pe.OPTIONAL_HEADER.CheckSum\n",
    "    res['Subsystem'] = pe.OPTIONAL_HEADER.Subsystem\n",
    "    res['DllCharacteristics'] = pe.OPTIONAL_HEADER.DllCharacteristics\n",
    "    res['SizeOfStackReserve'] = pe.OPTIONAL_HEADER.SizeOfStackReserve\n",
    "    res['SizeOfStackCommit'] = pe.OPTIONAL_HEADER.SizeOfStackCommit\n",
    "    res['SizeOfHeapReserve'] = pe.OPTIONAL_HEADER.SizeOfHeapReserve\n",
    "    res['SizeOfHeapCommit'] = pe.OPTIONAL_HEADER.SizeOfHeapCommit\n",
    "    res['LoaderFlags'] = pe.OPTIONAL_HEADER.LoaderFlags\n",
    "    res['NumberOfRvaAndSizes'] = pe.OPTIONAL_HEADER.NumberOfRvaAndSizes\n",
    "\n",
    "    # Sections\n",
    "    res['SectionsNb'] = len(pe.sections)\n",
    "    entropy = map(lambda x:x.get_entropy(), pe.sections)\n",
    "    res['SectionsMeanEntropy'] = sum(entropy)/float(len(entropy))\n",
    "    res['SectionsMinEntropy'] = min(entropy)\n",
    "    res['SectionsMaxEntropy'] = max(entropy)\n",
    "    raw_sizes = map(lambda x:x.SizeOfRawData, pe.sections)\n",
    "    res['SectionsMeanRawsize'] = sum(raw_sizes)/float(len(raw_sizes))\n",
    "    res['SectionsMinRawsize'] = min(raw_sizes)\n",
    "    res['SectionsMaxRawsize'] = max(raw_sizes)\n",
    "    virtual_sizes = map(lambda x:x.Misc_VirtualSize, pe.sections)\n",
    "    res['SectionsMeanVirtualsize'] = sum(virtual_sizes)/float(len(virtual_sizes))\n",
    "    res['SectionsMinVirtualsize'] = min(virtual_sizes)\n",
    "    res['SectionMaxVirtualsize'] = max(virtual_sizes)\n",
    "\n",
    "    #Imports\n",
    "    try:\n",
    "        res['ImportsNbDLL'] = len(pe.DIRECTORY_ENTRY_IMPORT)\n",
    "        imports = sum([x.imports for x in pe.DIRECTORY_ENTRY_IMPORT], [])\n",
    "        res['ImportsNb'] = len(imports)\n",
    "        res['ImportsNbOrdinal'] = len(filter(lambda x:x.name is None, imports))\n",
    "    except AttributeError:\n",
    "        res['ImportsNbDLL'] = 0\n",
    "        res['ImportsNb'] = 0\n",
    "        res['ImportsNbOrdinal'] = 0\n",
    "\n",
    "    #Exports\n",
    "    try:\n",
    "        res['ExportNb'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)\n",
    "    except AttributeError:\n",
    "        # No export\n",
    "        res['ExportNb'] = 0\n",
    "    #Resources\n",
    "    resources= get_resources(pe)\n",
    "    res['ResourcesNb'] = len(resources)\n",
    "    if len(resources)> 0:\n",
    "        entropy = map(lambda x:x[0], resources)\n",
    "        res['ResourcesMeanEntropy'] = sum(entropy)/float(len(entropy))\n",
    "        res['ResourcesMinEntropy'] = min(entropy)\n",
    "        res['ResourcesMaxEntropy'] = max(entropy)\n",
    "        sizes = map(lambda x:x[1], resources)\n",
    "        res['ResourcesMeanSize'] = sum(sizes)/float(len(sizes))\n",
    "        res['ResourcesMinSize'] = min(sizes)\n",
    "        res['ResourcesMaxSize'] = max(sizes)\n",
    "    else:\n",
    "        res['ResourcesNb'] = 0\n",
    "        res['ResourcesMeanEntropy'] = 0\n",
    "        res['ResourcesMinEntropy'] = 0\n",
    "        res['ResourcesMaxEntropy'] = 0\n",
    "        res['ResourcesMeanSize'] = 0\n",
    "        res['ResourcesMinSize'] = 0\n",
    "        res['ResourcesMaxSize'] = 0\n",
    "\n",
    "    # Load configuration size\n",
    "    try:\n",
    "        res['LoadConfigurationSize'] = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct.Size\n",
    "    except AttributeError:\n",
    "        res['LoadConfigurationSize'] = 0\n",
    "\n",
    "\n",
    "    # Version configuration size\n",
    "    try:\n",
    "        version_infos = get_version_info(pe)\n",
    "        res['VersionInformationSize'] = len(version_infos.keys())\n",
    "    except AttributeError:\n",
    "        res['VersionInformationSize'] = 0\n",
    "    return res\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "\t\n",
    "    clf = joblib.load('classifier/classifier.pkl')\n",
    "    features = pickle.loads(open(os.path.join('classifier/features.pkl'),'r').read())\n",
    "    data = extract_infos(sys.argv[1])\n",
    "    pe_features = map(lambda x:data[x], features)\n",
    "\n",
    "    res= clf.predict([pe_features])[0]    \n",
    "    print ('The file %s is %s' % (os.path.basename(sys.argv[1]),['malicious', 'legitimate'][res]))\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run the program to test the file - Skype.exe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The file Skype.exe is legitimate\n"
     ]
    }
   ],
   "source": [
    "%run malware_test.py \"/home/surajr/Downloads/Skype.exe\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test for the malicious file, an application has been downloaded from malwr.com"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The file BCN12ui49823.exe is malicious\n"
     ]
    }
   ],
   "source": [
    "%run malware_test.py \"/home/surajr/Downloads/BCN12ui49823.exe\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [Root]",
   "language": "python",
   "name": "Python [Root]"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}