CVE Tensorflow AI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "306d52aa",
   "metadata": {},
   "source": [
    "## CVE AI\n",
    "#### Classify CVE Security Vulnerabilities with Tensorflow\n",
    "\n",
    "This Jupyter Notebook will prepare and train one or multiple datasets which can be downloaded from the [NVD Website](https://nvd.nist.gov/vuln/data-feeds#JSON_FEED).\n",
    "\n",
    "The model will be trained with the textual description of each vulnerability and the corresponding severity High, Medium or Low.\n",
    "\n",
    "#### Training the model\n",
    "At the moment there are 19 datasets available (from 2002 until 2021). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "de832303",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import shutil\n",
    "import json\n",
    "import random\n",
    "\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import losses\n",
    "from tensorflow.keras import preprocessing\n",
    "from tensorflow.keras.layers.experimental.preprocessing import TextVectorization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ce16e436",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processing nvdcve-1.1-2019.json\n",
      "processing nvdcve-1.1-2020.json\n",
      "processing nvdcve-1.1-2017.json\n",
      "processing nvdcve-1.1-2018.json\n",
      "processing nvdcve-1.1-2021.json\n",
      "processing nvdcve-1.1-2002.json\n",
      "processing nvdcve-1.1-2003.json\n",
      "processing nvdcve-1.1-2004.json\n",
      "processing nvdcve-1.1-2005.json\n",
      "processing nvdcve-1.1-2006.json\n",
      "processing nvdcve-1.1-2007.json\n",
      "processing nvdcve-1.1-2008.json\n",
      "processing nvdcve-1.1-2009.json\n",
      "processing nvdcve-1.1-2010.json\n",
      "processing nvdcve-1.1-2011.json\n",
      "processing nvdcve-1.1-2012.json\n",
      "processing nvdcve-1.1-2013.json\n",
      "processing nvdcve-1.1-2014.json\n",
      "processing nvdcve-1.1-2015.json\n",
      "processing nvdcve-1.1-2016.json\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "# Read the dataset into memory\n",
    "\n",
    "# Multiple datasets from https://nvd.nist.gov/vuln/data-feeds#JSON_FEED can be read from the NVD_DATASET directory\n",
    "\n",
    "# Example directory structure:\n",
    "# NVD_DATASET\n",
    "#    - nvdcve-1.1-2020.json\n",
    "#    - nvdcve-1.1-2019.json\n",
    "\n",
    "\n",
    "raw_dataset_directory=\"./NVD_DATASET/\"\n",
    "directory_contents=os.listdir(raw_dataset_directory)\n",
    "\n",
    "nvd_dataset={}\n",
    "\n",
    "for filename in directory_contents:\n",
    "    if filename[-5:] == \".json\":\n",
    "        print(F\"processing {filename}\")\n",
    "        fp=open(F\"{raw_dataset_directory}/{filename}\", \"r\", encoding=\"utf-8\")\n",
    "        nvd_dataset[filename]=json.load(fp)\n",
    "        fp.close()\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4efdd105",
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:\n",
    "    print(F\"sample: {nvd['nvdcve-1.1-2019.json']['CVE_Items'][0]}\")\n",
    "    print(F\"\\n# of items: {len(nvd['nvdcve-1.1-2019.json']['CVE_Items'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0f623795",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "passed: 163375\n",
      "failed: 9990\n"
     ]
    }
   ],
   "source": [
    "# Extract the relevant data from the dataset\n",
    "\n",
    "# The textual description as well as the severity (HIGH, MEDIUM, or LOW) will be extracted from the dataset\n",
    "\n",
    "# Unfortunately the data structure of the JSON data is inconsistent\n",
    "# The statistics can be used to improve the data extraction\n",
    "\n",
    "passed = 0\n",
    "failed = 0\n",
    "item_pairs=[]\n",
    "\n",
    "for year in nvd_dataset.keys():\n",
    "    for item in nvd_dataset[year][\"CVE_Items\"]:\n",
    "        try:\n",
    "            severity = item[\"impact\"][\"baseMetricV2\"][\"severity\"] # severity\n",
    "            description = item[\"cve\"][\"description\"][\"description_data\"][0][\"value\"]\n",
    "            item_pairs.append((severity, description))\n",
    "            passed+=1\n",
    "        except:\n",
    "            failed+=1\n",
    "print(F\"passed: {passed}\\nfailed: {failed}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ea88e235",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "items available per severity class [52809, 93965, 16601]\n",
      "item limit: 16601\n"
     ]
    }
   ],
   "source": [
    "# Create balanced train and test datasets\n",
    "\n",
    "\n",
    "# sort\n",
    "high_severity_items=[i for i in item_pairs if i[0] == \"HIGH\"]\n",
    "medium_severity_items=[i for i in item_pairs if i[0] == \"MEDIUM\"]\n",
    "low_severity_items=[i for i in item_pairs if i[0] == \"LOW\"]\n",
    "\n",
    "# find out maximum possible item count for each severity class\n",
    "length_arr=[len(high_severity_items), len(medium_severity_items), len(low_severity_items)]\n",
    "item_num_limit=min(length_arr)\n",
    "print(F\"items available per severity class {length_arr}\")\n",
    "print(\"item limit:\",item_num_limit)\n",
    "\n",
    "# balance datasets\n",
    "low_severity_items=[random.choice(low_severity_items)[1] for i in range(item_num_limit)]\n",
    "medium_severity_items=[random.choice(medium_severity_items)[1] for i in range(item_num_limit)]\n",
    "high_severity_items=[random.choice(high_severity_items)[1] for i in range(item_num_limit)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3b5f76ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "directory NVD_PROCESSED exists\n",
      "files in train: 24900\n",
      "files in test: 24900\n",
      "LOW severity: 16601 items\n",
      "MEDIUM severity: 16601 items\n",
      "HIGH severity: 16601 items\n"
     ]
    }
   ],
   "source": [
    "# The datasets will be written to disk\n",
    "\n",
    "\n",
    "severity_categories=[\"LOW\", \"MEDIUM\", \"HIGH\"]\n",
    "nvd_prepared_dir=\"NVD_PROCESSED\"\n",
    "\n",
    "def create_directory_structure():\n",
    "    if os.path.exists(nvd_prepared_dir):\n",
    "        print(F\"directory {nvd_prepared_dir} exists\")\n",
    "        return\n",
    "    os.makedirs(nvd_prepared_dir)\n",
    "    os.makedirs(nvd_prepared_dir+\"/train\")\n",
    "    os.makedirs(nvd_prepared_dir+\"/test\")\n",
    "    for folder in severity_categories:\n",
    "        os.makedirs(nvd_prepared_dir+\"/train/\"+str(folder))\n",
    "        os.makedirs(nvd_prepared_dir+\"/test/\"+str(folder))\n",
    "\n",
    "def write_data_to_disk():\n",
    "    reportSev={sev:0 for sev in severity_categories}\n",
    "    reportDir={\"test\":0, \"train\":0}\n",
    "    for dataset,sev in [(high_severity_items,\"HIGH\"), (medium_severity_items,\"MEDIUM\"), (low_severity_items,\"LOW\")]:\n",
    "        filename=0\n",
    "        for subdir in [\"train\", \"test\"]:\n",
    "            datasetSub=[i for i in dataset[:len(dataset)//2]]\n",
    "            for desc in datasetSub:\n",
    "                filename+=1\n",
    "                fp=open(F\"{nvd_prepared_dir}/{subdir}/{sev}/{filename}.txt\", \"w\", encoding=\"utf-8\")\n",
    "                x=fp.write(desc.lower())\n",
    "                fp.close()\n",
    "            datasetSub=dataset[len(dataset)//2:]\n",
    "            reportDir[subdir]+=len(dataset)//2\n",
    "        reportSev[sev]+=len(dataset)\n",
    "    print(F\"files in train: {reportDir['train']}\\nfiles in test: {reportDir['test']}\")\n",
    "    for key in reportSev.keys():\n",
    "        print(F\"{key} severity: {reportSev[key]} items\")\n",
    "        \n",
    "def save_datasets():\n",
    "    create_directory_structure()\n",
    "    write_data_to_disk()\n",
    "\n",
    "if True: save_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9a94b8c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Severity categories:\n",
      "0 - LOW\n",
      "1 - MEDIUM\n",
      "2 - HIGH\n",
      "Found 24900 files belonging to 3 classes.\n",
      "Using 19920 files for training.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-11-03 20:18:51.164972: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:51.209882: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:51.210304: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:51.211459: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2021-11-03 20:18:51.211894: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:51.212298: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:51.212656: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:52.003603: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:52.004100: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:52.004511: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2021-11-03 20:18:52.004963: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1449 MB memory:  -> device: 0, name: NVIDIA GeForce MX450, pci bus id: 0000:01:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 24900 files belonging to 3 classes.\n",
      "Using 4980 files for validation.\n",
      "Found 41280 files belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Create Tensorflow training, validation and test datasets\n",
    "\n",
    "# 80% Training\n",
    "# 20% Validation\n",
    "\n",
    "severity_categories=[\"LOW\", \"MEDIUM\", \"HIGH\"]\n",
    "print(\"\\n\\nSeverity categories:\")\n",
    "for sevClass in range(len(severity_categories)):\n",
    "    print(F\"{sevClass} - {severity_categories[sevClass]}\")\n",
    "\n",
    "batch_size=32\n",
    "seed=42\n",
    "\n",
    "raw_train_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    F\"{nvd_prepared_dir}/train\",\n",
    "    batch_size=batch_size,\n",
    "    validation_split=0.2, # 80% will be used for training, 20% will be used for validation\n",
    "    subset=\"training\",\n",
    "    seed=seed)\n",
    "\n",
    "raw_val_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    F\"{nvd_prepared_dir}/train\",\n",
    "    batch_size=batch_size,\n",
    "    validation_split=0.2,\n",
    "    subset=\"validation\",\n",
    "    seed=seed)\n",
    "\n",
    "raw_test_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    F\"{nvd_prepared_dir}/test\", \n",
    "    batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fda76987",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-11-03 20:18:58.276517: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)\n"
     ]
    }
   ],
   "source": [
    "# Preprocess text to vector\n",
    "\n",
    "max_features = 100000\n",
    "embedding_dim = 512\n",
    "\n",
    "# Calculate the average length of all texts\n",
    "def avgLen(x): return len(x[1])\n",
    "item_pairs_len=list(map(avgLen, item_pairs))\n",
    "avgLen = int(sum(item_pairs_len)/len(item_pairs_len))\n",
    "sequence_length = avgLen\n",
    "\n",
    "vectorize_layer = TextVectorization(\n",
    "    # standardize=custom_standardization,\n",
    "    max_tokens=max_features,\n",
    "    output_mode=\"int\", # each token has a corresponding number\n",
    "    output_sequence_length=sequence_length)\n",
    "\n",
    "# adapt: map strings (words) to integers\n",
    "# Make a text-only dataset (without labels), then call adapt\n",
    "train_text = raw_train_ds.map(lambda x, y: x)\n",
    "vectorize_layer.adapt(train_text)\n",
    "\n",
    "# check result, preprocess some examples\n",
    "def vectorize_text(text, label):\n",
    "    text = tf.expand_dims(text, -1)\n",
    "    return vectorize_layer(text), label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a650c023",
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply to all sets, train, val, test must be the same\n",
    "train_ds = raw_train_ds.map(vectorize_text)\n",
    "val_ds = raw_val_ds.map(vectorize_text)\n",
    "test_ds = raw_test_ds.map(vectorize_text)\n",
    "\n",
    "# tuning performance\n",
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "\n",
    "train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5e4a2d9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (None, None, 512)         51200512  \n",
      "_________________________________________________________________\n",
      "dropout (Dropout)            (None, None, 512)         0         \n",
      "_________________________________________________________________\n",
      "global_average_pooling1d (Gl (None, 512)               0         \n",
      "_________________________________________________________________\n",
      "dropout_1 (Dropout)          (None, 512)               0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 128)               65664     \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 64)                8256      \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 3)                 195       \n",
      "=================================================================\n",
      "Total params: 51,274,627\n",
      "Trainable params: 51,274,627\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# compile the model\n",
    "model = tf.keras.Sequential([\n",
    "    layers.Embedding(max_features + 1, embedding_dim), # embedding layer creates an efficient, dense representation in which similar words have a similar encoding, improves through training\n",
    "    layers.Dropout(0.2),\n",
    "    layers.GlobalAveragePooling1D(), # reduces the length of each input vector to the average sequence length of all vectors\n",
    "    layers.Dropout(0.2),\n",
    "    layers.Dense(128, activation=\"relu\"),\n",
    "    layers.Dense(64, activation=\"relu\"),\n",
    "    layers.Dense(3)])\n",
    "\n",
    "model.summary()\n",
    "\n",
    "model.compile(\n",
    "    loss=losses.SparseCategoricalCrossentropy(from_logits=True), \n",
    "    optimizer=\"adam\", \n",
    "    metrics=[\"categorical_accuracy\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83aa26ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n"
     ]
    }
   ],
   "source": [
    "history = model.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    "    epochs=10)\n",
    "\n",
    "# evaluate the model\n",
    "loss, accuracy = model.evaluate(test_ds)\n",
    "print(\"Loss: \", loss)\n",
    "print(\"Accuracy: \", accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0ec9570",
   "metadata": {},
   "outputs": [],
   "source": [
    "export_model = tf.keras.Sequential([\n",
    "    vectorize_layer,\n",
    "    model,\n",
    "    layers.Activation(\"sigmoid\")])\n",
    "\n",
    "export_model.compile(\n",
    "    loss=losses.BinaryCrossentropy(from_logits=False), optimizer=\"adam\", metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9376e60",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(\"nvd_sev.tensorflow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "383e1700",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.42357512953367876 4632\n"
     ]
    }
   ],
   "source": [
    "# manually confirm accuracy\n",
    "highPred=export_model.predict(high_severity_items)\n",
    "medPred=export_model.predict(medium_severity_items)\n",
    "lowPred=export_model.predict(low_severity_items)\n",
    "numItems=len(highPred)\n",
    "\n",
    "stats=[0,0,0]\n",
    "for i in highPred:\n",
    "    if i[0] == max(i):\n",
    "        stats[0]+=1\n",
    "\n",
    "for i in medPred:\n",
    "    if i[1] == max(i):\n",
    "        stats[1]+=1\n",
    "\n",
    "for i in lowPred:\n",
    "    if i[2] == max(i):\n",
    "        stats[2]+=1\n",
    "\n",
    "avg_p=sum(stats)/3/numItems\n",
    "print(avg_p,numItems)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58676016",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}