BRON_cwe_capec.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:55.334570Z",
     "iopub.status.busy": "2022-05-29T20:23:55.325239Z",
     "iopub.status.idle": "2022-05-29T20:23:57.917412Z",
     "shell.execute_reply": "2022-05-29T20:23:57.915900Z"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, roc_auc_score, f1_score, confusion_matrix\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from transformers import AutoModel, AutoTokenizer, AutoModelForMaskedLM\n",
    "import json\n",
    "import spacy\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:57.942850Z",
     "iopub.status.busy": "2022-05-29T20:23:57.927934Z",
     "iopub.status.idle": "2022-05-29T20:23:57.994046Z",
     "shell.execute_reply": "2022-05-29T20:23:57.992797Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(\"w_dict.json\", \"r\") as f:\n",
    "    w_dict = json.load(f)\n",
    "with open(\"ap_dict.json\", \"r\") as f:\n",
    "    ap_dict = json.load(f)\n",
    "with open(\"technique_dict.json\", \"r\") as f:\n",
    "    technique_dict = json.load(f)\n",
    "with open(\"tactic_dict.json\", \"r\") as f:\n",
    "    tactic_dict = json.load(f)\n",
    "with open(\"cwe_names.json\", \"r\") as f:\n",
    "    cwe_names = json.load(f)\n",
    "with open(\"ap_names.json\", \"r\") as f:\n",
    "    ap_names = json.load(f)\n",
    "with open(\"technique_names.json\", \"r\") as f:\n",
    "    technique_names = json.load(f)\n",
    "with open(\"tactic_names.json\", \"r\") as f:\n",
    "    tactic_names = json.load(f)\n",
    "with open(\"cwe_short_descriptions.json\", \"r\") as f:\n",
    "    cwe_short_descriptions = json.load(f)\n",
    "with open(\"ap_short_descriptions.json\", \"r\") as f:\n",
    "    ap_short_descriptions = json.load(f)\n",
    "with open(\"technique_short_descriptions.json\", \"r\") as f:\n",
    "    technique_short_descriptions = json.load(f)\n",
    "with open(\"tactic_short_descriptions.json\", \"r\") as f:\n",
    "    tactic_short_descriptions = json.load(f)\n",
    "with open(\"cwe_descriptions.json\", \"r\") as f:\n",
    "    cwe_descriptions = json.load(f)\n",
    "with open(\"ap_descriptions.json\", \"r\") as f:\n",
    "    ap_descriptions = json.load(f)\n",
    "with open(\"technique_descriptions.json\", \"r\") as f:\n",
    "    technique_descriptions = json.load(f)\n",
    "with open(\"tactic_descriptions.json\", \"r\") as f:\n",
    "    tactic_descriptions = json.load(f)\n",
    "\n",
    "with open(\"ap_mitigation_descriptions.json\", \"r\") as f:\n",
    "    ap_mitigation_descriptions = json.load(f)\n",
    "with open(\"cwe_mitigation_descriptions.json\", \"r\") as f:\n",
    "    cwe_mitigation_descriptions = json.load(f)\n",
    "with open(\"tech_mitigation_names.json\", \"r\") as f:\n",
    "    tech_mitigation_names = json.load(f)\n",
    "with open(\"tech_detection_names.json\", \"r\") as f:\n",
    "    tech_detection_names = json.load(f)\n",
    "with open(\"ap_detection_descriptions.json\", \"r\") as f:\n",
    "    ap_detection_descriptions = json.load(f)\n",
    "with open(\"cwe_detection_descriptions.json\", \"r\") as f:\n",
    "    cwe_detection_descriptions = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:58.005166Z",
     "iopub.status.busy": "2022-05-29T20:23:58.003955Z",
     "iopub.status.idle": "2022-05-29T20:23:58.042199Z",
     "shell.execute_reply": "2022-05-29T20:23:58.043199Z"
    }
   },
   "outputs": [],
   "source": [
    "f = open(\"cwe_mitigation_ids_temp.json\")\n",
    "w_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"capec_mitigation_temp.json\")\n",
    "ap_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"technique_mitigation_temp.json\")\n",
    "technique_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"technique_detection_temp.json\")\n",
    "technique_detection = json.load(f)\n",
    "\n",
    "f = open(\"capec_detection_temp.json\")\n",
    "ap_detection = json.load(f)\n",
    "\n",
    "f = open(\"cwe_detection_temp.json\")\n",
    "w_detection = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:58.051702Z",
     "iopub.status.busy": "2022-05-29T20:23:58.050464Z",
     "iopub.status.idle": "2022-05-29T20:23:59.754321Z",
     "shell.execute_reply": "2022-05-29T20:23:59.755325Z"
    }
   },
   "outputs": [],
   "source": [
    "f = open(\"cwe_cve.json\")\n",
    "cwe_cve = json.load(f)\n",
    "\n",
    "f = open(\"cve.json\")\n",
    "vulnerabilities = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:59.797534Z",
     "iopub.status.busy": "2022-05-29T20:23:59.766637Z",
     "iopub.status.idle": "2022-05-29T20:23:59.882090Z",
     "shell.execute_reply": "2022-05-29T20:23:59.881016Z"
    }
   },
   "outputs": [],
   "source": [
    "v_dict = {}\n",
    "for cve in vulnerabilities:\n",
    "    if cve[\"original_id\"].split(\"-\")[1] == \"2021\":\n",
    "        v_dict[cve[\"_id\"]] = {\"description\": cve[\"metadata\"][\"description\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:23:59.912797Z",
     "iopub.status.busy": "2022-05-29T20:23:59.897490Z",
     "iopub.status.idle": "2022-05-29T20:24:00.177683Z",
     "shell.execute_reply": "2022-05-29T20:24:00.178940Z"
    }
   },
   "outputs": [],
   "source": [
    "# add CVEs to w_dict\n",
    "for cwe in w_dict:\n",
    "    w_dict[cwe][\"cve\"] = []\n",
    "count = 0\n",
    "cve_descriptions = set()\n",
    "\n",
    "for link in cwe_cve:\n",
    "    cwe = link[\"_from\"]\n",
    "    cve = link[\"_to\"]\n",
    "    if cve in v_dict:\n",
    "        cve_descriptions.add(v_dict[cve][\"description\"])\n",
    "        if cve not in w_dict[cwe][\"cve\"]:\n",
    "            w_dict[cwe][\"cve\"].append(cve)\n",
    "\n",
    "cve_descriptions = list(cve_descriptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:00.225247Z",
     "iopub.status.busy": "2022-05-29T20:24:00.220141Z",
     "iopub.status.idle": "2022-05-29T20:24:00.258113Z",
     "shell.execute_reply": "2022-05-29T20:24:00.257116Z"
    }
   },
   "outputs": [],
   "source": [
    "# consistent set of negative (cwe, capec) examples\n",
    "negative_example_ids = []\n",
    "\n",
    "random.seed(0)\n",
    "\n",
    "positive_example_num = 0\n",
    "\n",
    "for ap in ap_dict:\n",
    "    for cwe in ap_dict[ap][\"cwes\"]:\n",
    "        positive_example_num += 1\n",
    "\n",
    "for i in range(positive_example_num):\n",
    "    random_ap = random.choice(list(ap_dict))\n",
    "    random_cwe = random.choice(list(w_dict))\n",
    "    while (\n",
    "        random_cwe in ap_dict[random_ap][\"cwes\"]\n",
    "        or (random_ap, random_cwe) in negative_example_ids\n",
    "    ):\n",
    "        random_ap = random.choice(list(ap_dict))\n",
    "        random_cwe = random.choice(list(w_dict))\n",
    "    ap = random_ap\n",
    "    cwe = random_cwe\n",
    "    negative_example_ids.append((ap, cwe))\n",
    "\n",
    "negative_example_json = []\n",
    "for example in negative_example_ids:\n",
    "    negative_example_json.append({\"ap\": example[0], \"cwe\": example[1]})\n",
    "\n",
    "with open(\"negative_examples_100_cwe_capec.json\", \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(negative_example_json, f, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:00.267785Z",
     "iopub.status.busy": "2022-05-29T20:24:00.264278Z",
     "iopub.status.idle": "2022-05-29T20:24:00.273152Z",
     "shell.execute_reply": "2022-05-29T20:24:00.272394Z"
    }
   },
   "outputs": [],
   "source": [
    "negative_example_ids = []\n",
    "negative_example_json = json.load(open(\"negative_examples_100_cwe_capec.json\"))\n",
    "for example in negative_example_json:\n",
    "    negative_example_ids.append((example[\"ap\"], example[\"cwe\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:00.283472Z",
     "iopub.status.busy": "2022-05-29T20:24:00.282260Z",
     "iopub.status.idle": "2022-05-29T20:24:00.286514Z",
     "shell.execute_reply": "2022-05-29T20:24:00.287184Z"
    }
   },
   "outputs": [],
   "source": [
    "positive_example_ids = []\n",
    "for ap in ap_dict:\n",
    "    for cwe in ap_dict[ap][\"cwes\"]:\n",
    "        positive_example_ids.append((ap, cwe))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:00.294272Z",
     "iopub.status.busy": "2022-05-29T20:24:00.293314Z",
     "iopub.status.idle": "2022-05-29T20:24:00.298148Z",
     "shell.execute_reply": "2022-05-29T20:24:00.299015Z"
    }
   },
   "outputs": [],
   "source": [
    "example_ids = positive_example_ids + negative_example_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:00.307805Z",
     "iopub.status.busy": "2022-05-29T20:24:00.306621Z",
     "iopub.status.idle": "2022-05-29T20:24:05.122394Z",
     "shell.execute_reply": "2022-05-29T20:24:05.123091Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "device = \"cpu\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "pretrained_model = AutoModel.from_pretrained(\"bert-base-uncased\").to(device)\n",
    "\n",
    "model_path = \"bert_base\"\n",
    "finetuned_model = AutoModelForMaskedLM.from_pretrained(model_path).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:05.128861Z",
     "iopub.status.busy": "2022-05-29T20:24:05.128002Z",
     "iopub.status.idle": "2022-05-29T20:24:07.217702Z",
     "shell.execute_reply": "2022-05-29T20:24:07.216166Z"
    }
   },
   "outputs": [],
   "source": [
    "encode = spacy.load(\"en_core_web_lg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:07.227871Z",
     "iopub.status.busy": "2022-05-29T20:24:07.226570Z",
     "iopub.status.idle": "2022-05-29T20:24:07.231225Z",
     "shell.execute_reply": "2022-05-29T20:24:07.232186Z"
    }
   },
   "outputs": [],
   "source": [
    "results = {\n",
    "    \"1\": {},\n",
    "    \"2\": {},\n",
    "    \"3\": {},\n",
    "    \"4\": {},\n",
    "    \"5\": {},\n",
    "    \"6\": {},\n",
    "    \"7\": {},\n",
    "    \"8\": {},\n",
    "    \"9\": {},\n",
    "    \"10\": {},\n",
    "    \"11\": {},\n",
    "    \"12\": {},\n",
    "    \"13\": {},\n",
    "    \"14\": {},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:07.253024Z",
     "iopub.status.busy": "2022-05-29T20:24:07.251204Z",
     "iopub.status.idle": "2022-05-29T20:24:08.340145Z",
     "shell.execute_reply": "2022-05-29T20:24:08.341369Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CountVectorizer()"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ap_name_vectorizer = CountVectorizer()\n",
    "ap_name_vectorizer.fit(ap_names)\n",
    "\n",
    "technique_name_vectorizer = CountVectorizer()\n",
    "technique_name_vectorizer.fit(technique_names)\n",
    "\n",
    "cwe_name_vectorizer = CountVectorizer()\n",
    "cwe_name_vectorizer.fit(cwe_names)\n",
    "\n",
    "tactic_name_vectorizer = CountVectorizer()\n",
    "tactic_name_vectorizer.fit(tactic_names)\n",
    "\n",
    "ap_mitigation_vectorizer = CountVectorizer()\n",
    "cwe_mitigation_vectorizer = CountVectorizer()\n",
    "tech_mitigation_vectorizer = CountVectorizer()\n",
    "\n",
    "ap_mitigation_vectorizer.fit(ap_mitigation_descriptions)\n",
    "cwe_mitigation_vectorizer.fit(cwe_mitigation_descriptions)\n",
    "tech_mitigation_vectorizer.fit(tech_mitigation_names)\n",
    "\n",
    "ap_detection_vectorizer = CountVectorizer()\n",
    "cwe_detection_vectorizer = CountVectorizer()\n",
    "tech_detection_vectorizer = CountVectorizer()\n",
    "\n",
    "ap_detection_vectorizer.fit(ap_detection_descriptions)\n",
    "cwe_detection_vectorizer.fit(cwe_detection_descriptions)\n",
    "tech_detection_vectorizer.fit(tech_detection_names)\n",
    "\n",
    "ap_description_vectorizer = CountVectorizer()\n",
    "ap_description_vectorizer.fit(ap_descriptions)\n",
    "\n",
    "cwe_description_vectorizer = CountVectorizer()\n",
    "cwe_description_vectorizer.fit(cwe_descriptions)\n",
    "\n",
    "technique_description_vectorizer = CountVectorizer()\n",
    "technique_description_vectorizer.fit(technique_descriptions)\n",
    "\n",
    "tactic_description_vectorizer = CountVectorizer()\n",
    "tactic_description_vectorizer.fit(tactic_descriptions)\n",
    "\n",
    "ap_short_description_vectorizer = CountVectorizer()\n",
    "ap_short_description_vectorizer.fit(ap_short_descriptions)\n",
    "\n",
    "cwe_short_description_vectorizer = CountVectorizer()\n",
    "cwe_short_description_vectorizer.fit(cwe_short_descriptions)\n",
    "\n",
    "technique_short_description_vectorizer = CountVectorizer()\n",
    "technique_short_description_vectorizer.fit(technique_short_descriptions)\n",
    "\n",
    "tactic_short_description_vectorizer = CountVectorizer()\n",
    "tactic_short_description_vectorizer.fit(tactic_short_descriptions)\n",
    "\n",
    "cve_description_vectorizer = CountVectorizer()\n",
    "cve_description_vectorizer.fit(cve_descriptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:08.397332Z",
     "iopub.status.busy": "2022-05-29T20:24:08.353208Z",
     "iopub.status.idle": "2022-05-29T20:24:08.434335Z",
     "shell.execute_reply": "2022-05-29T20:24:08.435336Z"
    }
   },
   "outputs": [],
   "source": [
    "def vector_encoding(\n",
    "    encoding_type, text, vectorizer=None, bert_output_type=None, bert_finetuned=False\n",
    "):\n",
    "    if encoding_type == \"None\":\n",
    "        return text\n",
    "    elif encoding_type == \"BoW\":\n",
    "        return vectorizer_transform(text, vectorizer)\n",
    "    elif encoding_type == \"spaCy\":\n",
    "        return spaCy_vector(text)\n",
    "    elif encoding_type == \"BERT\":\n",
    "        if bert_finetuned:\n",
    "            model = finetuned_model\n",
    "        else:\n",
    "            model = pretrained_model\n",
    "\n",
    "        if bert_output_type == \"pooler_output\":\n",
    "            return get_pooler_output(model, text)\n",
    "        elif bert_output_type == \"hidden_state\":\n",
    "            return get_hidden_state(model, text)\n",
    "\n",
    "\n",
    "def vectorizer_transform(input_to_BoW, vectorizer):\n",
    "    return vectorizer.transform([input_to_BoW])[0].toarray().flatten()\n",
    "\n",
    "\n",
    "def spaCy_vector(text):\n",
    "    return encode(text).vector\n",
    "\n",
    "\n",
    "def get_pooler_output(model, text):\n",
    "    inputs = tokenizer(text.lower(), truncation=True, return_tensors=\"pt\").to(device)\n",
    "    outputs = model(**inputs)\n",
    "    pooled_output = outputs.pooler_output\n",
    "    return pooled_output.detach().cpu().numpy().flatten()\n",
    "\n",
    "\n",
    "def get_hidden_state(model, text):\n",
    "    inputs = tokenizer(text.lower(), truncation=True, return_tensors=\"pt\").to(device)\n",
    "    outputs = model(**inputs, output_hidden_states=True)\n",
    "    hidden_states = outputs.hidden_states\n",
    "    return hidden_states[-1][:, 0, :].detach().cpu().numpy().flatten()\n",
    "\n",
    "\n",
    "def append_data(\n",
    "    encoding_type, data_combo, ap, cwe, bert_output_type=None, bert_finetuned=False\n",
    "):\n",
    "    output = []\n",
    "    vectorizer = CountVectorizer()\n",
    "\n",
    "    ap_text = ap_names\n",
    "    technique_text = technique_names\n",
    "    cwe_text = cwe_names\n",
    "    tactic_text = tactic_names\n",
    "    cve_text = cve_descriptions\n",
    "\n",
    "    if data_combo == \"A0\":\n",
    "        vectorizer.fit(ap_text + cwe_text)\n",
    "\n",
    "    elif data_combo == \"A1\":\n",
    "        vectorizer.fit(ap_text + technique_text + cwe_text + tactic_text + cve_text)\n",
    "\n",
    "    elif data_combo == \"A1 + MI\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cve_text\n",
    "            + cwe_mitigation_descriptions\n",
    "            + ap_mitigation_descriptions\n",
    "            + tech_mitigation_names\n",
    "        )\n",
    "\n",
    "    elif data_combo == \"A1 + D\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cve_text\n",
    "            + cwe_detection_descriptions\n",
    "            + ap_detection_descriptions\n",
    "            + tech_detection_names\n",
    "        )\n",
    "\n",
    "    elif data_combo == \"A1 + MI + D\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cve_text\n",
    "            + cwe_mitigation_descriptions\n",
    "            + ap_mitigation_descriptions\n",
    "            + tech_mitigation_names\n",
    "            + cwe_detection_descriptions\n",
    "            + ap_detection_descriptions\n",
    "            + tech_detection_names\n",
    "        )\n",
    "\n",
    "    output.append(ap_dict[ap][\"name\"])\n",
    "    output.append(w_dict[cwe][\"name\"])\n",
    "\n",
    "    techniques = ap_dict[ap][\"techniques\"]\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        output.append(w_dict[cwe][\"name\"])\n",
    "\n",
    "        for technique in techniques:\n",
    "            output.append(technique_dict[technique][\"name\"])\n",
    "\n",
    "        for technique in techniques:\n",
    "            for tac in technique_dict[technique][\"tactics\"]:\n",
    "                output.append(tactic_dict[tac][\"name\"])\n",
    "\n",
    "        for cve in w_dict[cwe][\"cve\"]:\n",
    "            output.append(v_dict[cve][\"description\"])\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "            for cwe_mit in w_mitigation:\n",
    "                if mitigation == cwe_mit[\"_id\"]:\n",
    "                    output.append(cwe_mit[\"metadata\"][\"Description\"])\n",
    "\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            for ap_mit in ap_mitigation:\n",
    "                if mitigation == ap_mit[\"_id\"]:\n",
    "                    output.append(ap_mit[\"metadata\"])\n",
    "\n",
    "        for technique in techniques:\n",
    "            for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "                for tech_mit in technique_mitigation:\n",
    "                    if mitigation == tech_mit[\"_id\"]:\n",
    "                        output.append(tech_mit[\"name\"])\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        for detection in w_dict[cwe][\"detections\"]:\n",
    "            for cwe_det in w_detection:\n",
    "                if detection == cwe_det[\"_id\"]:\n",
    "                    output.append(cwe_det[\"metadata\"][\"Description\"])\n",
    "\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            for ap_det in ap_detection:\n",
    "                if detection == ap_det[\"_id\"]:\n",
    "                    output.append(ap_det[\"metadata\"])\n",
    "\n",
    "        for technique in techniques:\n",
    "            for detection in technique_dict[technique][\"detections\"]:\n",
    "                for tech_det in technique_mitigation:\n",
    "                    if detection == tech_det[\"_id\"]:\n",
    "                        output.append(tech_det[\"metadata\"])\n",
    "\n",
    "    output = \" \".join(output)\n",
    "    return vector_encoding(\n",
    "        encoding_type, output, vectorizer, bert_output_type, bert_finetuned\n",
    "    )\n",
    "\n",
    "\n",
    "def handle_data(\n",
    "    encoding_type, data_combo, ap, cwe, bert_output_type=None, bert_finetuned=False\n",
    "):\n",
    "    example = []\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            ap_dict[ap][\"name\"],\n",
    "            ap_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            w_dict[cwe][\"name\"],\n",
    "            cwe_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    techniques = ap_dict[ap][\"techniques\"]\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        techs = []\n",
    "        for technique in techniques:\n",
    "            techs.append(technique_dict[technique][\"name\"])\n",
    "        techs = \" \".join(techs)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                techs,\n",
    "                technique_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        tactics = []\n",
    "        for technique in techniques:\n",
    "            for tac in technique_dict[technique][\"tactics\"]:\n",
    "                tactics.append(tactic_dict[tac][\"name\"])\n",
    "        tactics = \" \".join(tactics)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tactics,\n",
    "                tactic_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        cves = []\n",
    "        for cve in w_dict[cwe][\"cve\"]:\n",
    "            cves.append(v_dict[cve][\"description\"])\n",
    "        cves = \" \".join(cves)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cves,\n",
    "                cve_description_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        cwe_mitigations = []\n",
    "        for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "            for cwe_mit in w_mitigation:\n",
    "                if mitigation == cwe_mit[\"_id\"]:\n",
    "                    cwe_mitigations.append(cwe_mit[\"metadata\"][\"Description\"])\n",
    "        cwe_mitigations = \" \".join(cwe_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cwe_mitigations,\n",
    "                cwe_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        capec_mitigations = []\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            for ap_mit in ap_mitigation:\n",
    "                if mitigation == ap_mit[\"_id\"]:\n",
    "                    capec_mitigations.append(ap_mit[\"metadata\"])\n",
    "        capec_mitigations = \" \".join(capec_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                capec_mitigations,\n",
    "                ap_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        tech_mitigations = []\n",
    "        for technique in techniques:\n",
    "            for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "                for tech_mit in technique_mitigation:\n",
    "                    if mitigation == tech_mit[\"_id\"]:\n",
    "                        tech_mitigations.append(tech_mit[\"name\"])\n",
    "        tech_mitigations = \" \".join(tech_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tech_mitigations,\n",
    "                tech_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        cwe_detections = []\n",
    "        for detection in w_dict[cwe][\"detections\"]:\n",
    "            for cwe_det in w_detection:\n",
    "                if detection == cwe_det[\"_id\"]:\n",
    "                    cwe_detections.append(cwe_det[\"metadata\"][\"Description\"])\n",
    "\n",
    "        cwe_detections = \" \".join(cwe_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cwe_detections,\n",
    "                cwe_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        capec_detections = []\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            for ap_det in ap_detection:\n",
    "                if detection == ap_det[\"_id\"]:\n",
    "                    capec_detections.append(ap_det[\"metadata\"])\n",
    "        capec_detections = \" \".join(capec_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                capec_detections,\n",
    "                ap_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        tech_detections = []\n",
    "        for technique in techniques:\n",
    "            for detection in technique_dict[technique][\"detections\"]:\n",
    "                for tech_det in technique_detection:\n",
    "                    if detection == tech_det[\"_id\"]:\n",
    "                        tech_detections.append(tech_det[\"name\"])\n",
    "        tech_detections = \" \".join(tech_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tech_detections,\n",
    "                tech_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    return np.hstack(example)\n",
    "\n",
    "\n",
    "def encode_data(\n",
    "    encoding_type, data_combo, ap, cwe, bert_output_type=None, bert_finetuned=False\n",
    "):\n",
    "    example = []\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            ap_dict[ap][\"name\"],\n",
    "            ap_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            w_dict[cwe][\"name\"],\n",
    "            cwe_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    techniques = ap_dict[ap][\"techniques\"]\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        for technique in techniques:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    technique_dict[technique][\"name\"],\n",
    "                    technique_name_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        for technique in techniques:\n",
    "            for tac in technique_dict[technique][\"tactics\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        tactic_dict[tac][\"name\"],\n",
    "                        tactic_name_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        for cve in w_dict[cwe][\"cve\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    v_dict[cve][\"description\"],\n",
    "                    cve_description_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "            for cwe_mit in w_mitigation:\n",
    "                if mitigation == cwe_mit[\"_id\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            cwe_mit[\"metadata\"][\"Description\"],\n",
    "                            cwe_mitigation_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            for ap_mit in ap_mitigation:\n",
    "                if mitigation == ap_mit[\"_id\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            ap_mit[\"metadata\"],\n",
    "                            ap_mitigation_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for technique in techniques:\n",
    "            for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "                for tech_mit in technique_mitigation:\n",
    "                    if mitigation == tech_mit[\"_id\"]:\n",
    "                        example.append(\n",
    "                            vector_encoding(\n",
    "                                encoding_type,\n",
    "                                tech_mit[\"name\"],\n",
    "                                tech_mitigation_vectorizer,\n",
    "                                bert_output_type,\n",
    "                                bert_finetuned,\n",
    "                            )\n",
    "                        )\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        for detection in w_dict[cwe][\"detections\"]:\n",
    "            for cwe_det in w_detection:\n",
    "                if detection == cwe_det[\"_id\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            cwe_det[\"metadata\"][\"Description\"],\n",
    "                            cwe_detection_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            for ap_det in ap_detection:\n",
    "                if detection == ap_det[\"_id\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            ap_det[\"metadata\"],\n",
    "                            ap_detection_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for technique in techniques:\n",
    "            for detection in technique_dict[technique][\"detections\"]:\n",
    "                for tech_det in technique_detection:\n",
    "                    if detection == tech_det[\"_id\"]:\n",
    "                        example.append(\n",
    "                            vector_encoding(\n",
    "                                encoding_type,\n",
    "                                tech_det[\"name\"],\n",
    "                                tech_detection_vectorizer,\n",
    "                                bert_output_type,\n",
    "                                bert_finetuned,\n",
    "                            )\n",
    "                        )\n",
    "\n",
    "    return np.hstack(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T20:24:08.479612Z",
     "iopub.status.busy": "2022-05-29T20:24:08.478673Z",
     "iopub.status.idle": "2022-05-29T23:43:03.603680Z",
     "shell.execute_reply": "2022-05-29T23:43:03.601263Z"
    }
   },
   "outputs": [],
   "source": [
    "encodings_to_use = [1, 2]\n",
    "\n",
    "max_length_BoW = 0\n",
    "max_length_spaCy = 0\n",
    "max_length_BERT = 0\n",
    "for data_combo in [\"A0\", \"A1\", \"A1 + MI\", \"A1 + D\", \"A1 + MI + D\"]:\n",
    "    for encoding in encodings_to_use:\n",
    "        accuracies = []\n",
    "        aucs = []\n",
    "        f1s = []\n",
    "        fps = []\n",
    "        labels = []\n",
    "        examples = []\n",
    "        for i, (ap, cwe) in enumerate(example_ids):\n",
    "            if encoding == 1:\n",
    "                examples.append(append_data(\"BoW\", data_combo, ap, cwe))\n",
    "\n",
    "            elif encoding == 2:\n",
    "                examples.append(handle_data(\"BoW\", data_combo, ap, cwe))\n",
    "\n",
    "            elif encoding == 3:\n",
    "                examples.append(encode_data(\"BoW\", data_combo, ap, cwe))\n",
    "                max_length_BoW = max(max_length_BoW, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 4:\n",
    "                examples.append(append_data(\"spaCy\", data_combo, ap, cwe))\n",
    "\n",
    "            elif encoding == 5:\n",
    "                examples.append(handle_data(\"spaCy\", data_combo, ap, cwe))\n",
    "\n",
    "            elif encoding == 6:\n",
    "                examples.append(encode_data(\"spaCy\", data_combo, ap, cwe))\n",
    "                max_length_spaCy = max(max_length_spaCy, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 7:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"pooler_output\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 8:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"pooler_output\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 9:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 10:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 11:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 12:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 13:\n",
    "                examples.append(\n",
    "                    handle_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 14:\n",
    "                examples.append(\n",
    "                    handle_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        cwe,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            if i < len(example_ids) / 2:\n",
    "                labels.append(1)\n",
    "            else:\n",
    "                labels.append(0)\n",
    "\n",
    "        if encoding == 3:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_BoW - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        if encoding == 6:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_spaCy - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        if encoding == 7 or encoding == 9 or encoding == 11:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_BERT - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        for i in range(100):\n",
    "            # training/testing classifier\n",
    "            X_train, X_test, y_train, y_test = train_test_split(\n",
    "                examples, labels, test_size=0.3, random_state=i\n",
    "            )\n",
    "            clf = RandomForestClassifier(random_state=i, class_weight={0: 1, 1: 1})\n",
    "            clf.fit(X_train, y_train)\n",
    "            y_pred = clf.predict(X_test)\n",
    "\n",
    "            confusion = confusion_matrix(y_test, y_pred)\n",
    "            fp_rate = confusion[0][1] / (confusion[0][1] + confusion[0][0])\n",
    "            fps.append(fp_rate)\n",
    "\n",
    "            accuracies.append(accuracy_score(y_test, y_pred))\n",
    "            aucs.append(roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1]))\n",
    "            f1s.append(f1_score(y_test, y_pred))\n",
    "\n",
    "        results[str(encoding)][data_combo] = {\n",
    "            \"fp\": fps,\n",
    "            \"acc\": accuracies,\n",
    "            \"auc\": aucs,\n",
    "            \"f1\": f1s,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-29T23:43:03.618045Z",
     "iopub.status.busy": "2022-05-29T23:43:03.616292Z",
     "iopub.status.idle": "2022-05-29T23:43:03.638593Z",
     "shell.execute_reply": "2022-05-29T23:43:03.637504Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(\"results_100_with_BoW_cwe_capec_with_cve_2021.json\", \"w\") as f:\n",
    "    json.dump(results, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "98b0a9b7b4eaaa670588a142fd0a9b87eaafe866f1db4228be72b4211d12040f"
  },
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('base': conda)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}