Model POC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import lightgbm as lgb\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import config\n",
    "import builder\n",
    "import transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fetch Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dataset is fetched into memory via the function _Reader_.\n",
    "It reads the features from every malware classes' _dump.json_ in their respective folders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "\u001b[36m[*] Initiated dataset read.\u001b[00m\n\u001b[34m[*] Reading dataset for backdoor.\u001b[00m\n\u001b[91m[+] Dataset fetch for backdoor complete.\u001b[00m\n\u001b[34m[*] Reading dataset for worm.\u001b[00m\n\u001b[91m[+] Dataset fetch for worm complete.\u001b[00m\n\u001b[34m[*] Reading dataset for trojan.\u001b[00m\n\u001b[91m[+] Dataset fetch for trojan complete.\u001b[00m\n\u001b[34m[*] Reading dataset for rootkit.\u001b[00m\n\u001b[91m[+] Dataset fetch for rootkit complete.\u001b[00m\n\u001b[34m[*] Reading dataset for virus.\u001b[00m\n\u001b[91m[+] Dataset fetch for virus complete.\u001b[00m\n\u001b[34m[*] Reading dataset for bot.\u001b[00m\n\u001b[91m[+] Dataset fetch for bot complete.\u001b[00m\n\u001b[34m[*] Reading dataset for ransomware.\u001b[00m\n\u001b[91m[+] Dataset fetch for ransomware complete.\u001b[00m\n\u001b[34m[*] Reading dataset for adware.\u001b[00m\n\u001b[91m[+] Dataset fetch for adware complete.\u001b[00m\n\u001b[34m[*] Reading dataset for downloader.\u001b[00m\n\u001b[91m[+] Dataset fetch for downloader complete.\u001b[00m\n\u001b[91m[+] Dataset loading complete.\u001b[00m\n"
    }
   ],
   "source": [
    "# Setup dataset for training.\n",
    "X = []\n",
    "Y = []\n",
    "dataset = builder.Reader()\n",
    "# Iterate over dataset through all specified classes.\n",
    "for typeClass in config.Classes:\n",
    "    X += dataset[typeClass]\n",
    "    # Append labels for all the elements fetched in the given class.\n",
    "    for i in range(len(dataset[typeClass])):\n",
    "        Y.append(list(dataset.keys()).index(typeClass))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Form the datasets for training and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "Shape of data:  (422, 2152)\nShape of labels:  (422,)\n"
    }
   ],
   "source": [
    "Data = np.array(X, dtype=np.float32)\n",
    "Labels = np.asarray(Y, dtype=np.float32)\n",
    "print(\"Shape of data: \", Data.shape)\n",
    "print(\"Shape of labels: \", Labels.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data into dataframes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "[0. 1. 2. 3. 4. 5. 6. 7. 8.]\n(422,)\n[0. 1. 2. 3. 4. 5. 6. 7. 8.]\n"
    }
   ],
   "source": [
    "# Load respective dataframes.\n",
    "X = pd.DataFrame(Data)\n",
    "Y = pd.DataFrame(Labels)\n",
    "\n",
    "# Make sure samples from all classes are present.\n",
    "print(Y[0].unique())\n",
    "\n",
    "# Segregate the data and labels from the same dataframe to prevent inconsistency.\n",
    "frames = [X, Y]\n",
    "DATA = pd.concat(frames, axis = 1)\n",
    "Y = DATA.iloc[:, -1]\n",
    "print(Y.shape)\n",
    "X = DATA.iloc[:, :-1]\n",
    "print(np.unique(Y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split the dataset for training and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "Training data: (337, 2152)\nTraining labels: (337,)\nTesting data: (85, 2152)\nTesting labels: (85,)\n"
    }
   ],
   "source": [
    "x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size = 0.2, random_state = 42)\n",
    "\n",
    "x_train = np.array(x_train)\n",
    "print(\"Training data:\", np.shape(x_train))\n",
    "\n",
    "y_train = np.array(y_train)\n",
    "print(\"Training labels:\", np.shape(y_train))\n",
    "\n",
    "x_test = np.array(x_test)\n",
    "print(\"Testing data:\", np.shape(x_test))\n",
    "\n",
    "y_test = np.array(y_test)\n",
    "print(\"Testing labels:\", np.shape(y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "Number of samples for training: 337\nNumber of samples for testing: 85\nNumber of features for each sample: 2152\n"
    }
   ],
   "source": [
    "print(\"Number of samples for training:\", np.shape(x_train)[0])\n",
    "print(\"Number of samples for testing:\", np.shape(x_test)[0])\n",
    "print(\"Number of features for each sample:\", np.shape(x_train)[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the parameters for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {}\n",
    "params[\"learning_rate\"] = 0.05\n",
    "params[\"boosting_type\"] = \"gbdt\"\n",
    "params[\"objective\"] = \"multiclass\"\n",
    "params[\"num_class\"] = len(config.Classes)\n",
    "params[\"metric\"] = \"multi_logloss\"\n",
    "params[\"sub_feature\"] = 0.3\n",
    "params[\"num_leaves\"] = 15\n",
    "params[\"min_data\"] = 95\n",
    "params[\"max_depth\"] = 15\n",
    "params[\"device\"] = \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup the dataset for training.\n",
    "d_train = lgb.Dataset(x_train, label = y_train)\n",
    "\n",
    "# Train the model based on the aforementioned dataset.\n",
    "clf = lgb.train(params, d_train, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check accuracy on test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.7376470588235295"
     },
     "metadata": {},
     "execution_count": 10
    }
   ],
   "source": [
    "y_pred = clf.predict(x_test)\n",
    "best_preds = [np.argmax(line) for line in y_pred]\n",
    "accuracy_score(y_test, best_preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "<lightgbm.basic.Booster at 0x7fe94a227d00>"
     },
     "metadata": {},
     "execution_count": 11
    }
   ],
   "source": [
    "clf.save_model(\"model.mdl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function to make prediction for a given input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Prediction(pe_file, mdlFile):\n",
    "    \n",
    "    predictor = lgb.Booster(model_file = mdlFile)    \n",
    "    # Fetch the feature vector for the PE.\n",
    "    transformed = transformer.PETransformer(peFile).vector\n",
    "    # Make prediction for the PE.\n",
    "    preds = predictor.predict(transformed.reshape(1, 2152))\n",
    "    # Gives the maximum value out of all the predicted labels.\n",
    "    return config.Classes[np.argmax(preds)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sample prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "trojan\n"
    }
   ],
   "source": [
    "peFile = \"dataset/trojan/04eacd2031de21c56ccec496e1b5ed68\"\n",
    "\n",
    "# Get the data from the file.\n",
    "data = open(peFile, \"rb\").read()\n",
    "# Predict the class of the file.\n",
    "print(Prediction(data, \"model.mdl\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ML",
   "language": "python",
   "name": "ml"
  },
  "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.3-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}