generate_link_probabilities_cwe_capec.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "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.ensemble import RandomForestClassifier\n",
    "from transformers import AutoModel, AutoTokenizer, AutoModelForMaskedLM\n",
    "import json\n",
    "import spacy\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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": 3,
   "metadata": {},
   "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(\"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": 4,
   "metadata": {},
   "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": 5,
   "metadata": {},
   "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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_ids = positive_example_ids + negative_example_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CountVectorizer()"
      ]
     },
     "execution_count": 7,
     "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "encode = spacy.load(\"en_core_web_lg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.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": 10,
   "metadata": {},
   "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\")\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\")\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",
    "\n",
    "    if data_combo == \"A0\":\n",
    "        vectorizer.fit(ap_text + technique_text)\n",
    "\n",
    "    elif data_combo == \"A1\":\n",
    "        vectorizer.fit(ap_text + technique_text + cwe_text + tactic_text)\n",
    "\n",
    "    elif data_combo == \"A1 + MI\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_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",
    "            + 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",
    "            + 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",
    "    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",
    "    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",
    "    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": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = []\n",
    "labels = []\n",
    "for i, (ap, cwe) in enumerate(example_ids):\n",
    "    # examples.append(append_data(\"BoW\", \"A0\", ap, cwe))\n",
    "    examples.append(handle_data(\"spaCy\", \"A0\", ap, cwe))\n",
    "    if i < len(example_ids) / 2:\n",
    "        labels.append(1)\n",
    "    else:\n",
    "        labels.append(0)\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    examples, labels, test_size=0.3, random_state=0\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomForestClassifier(class_weight={0: 1, 1: 1}, random_state=0)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clf = RandomForestClassifier(random_state=0, class_weight={0: 1, 1: 1})\n",
    "clf.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "filename = \"link_probabilities_model_spaCy_cwe_capec.pkl\"\n",
    "with open(filename, \"wb\") as f:\n",
    "    pickle.dump(clf, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather_text_data(ap, cwe, data_combo):\n",
    "    output = {}\n",
    "    output[\"AP\"] = ap_dict[ap][\"name\"]\n",
    "    output[\"CWE\"] = w_dict[cwe][\"name\"]\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        output[\"Techniques\"] = []\n",
    "        for technique in ap_dict[ap][\"techniques\"]:\n",
    "            output[\"Techniques\"].append(technique_dict[technique][\"name\"])\n",
    "        output[\"Tactics\"] = []\n",
    "        for technique in ap_dict[ap][\"techniques\"]:\n",
    "            for tac in technique_dict[technique][\"tactics\"]:\n",
    "                output[\"Tactics\"].append(tactic_dict[tac][\"name\"])\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'filename = \\'link_probabilities_model_BERT.pkl\\'\\nwith open(filename, \"rb\") as f:\\n    BERT_clf = pickle.load(f)'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# if classifier already trained\n",
    "import pickle\n",
    "\n",
    "filename = \"link_probabilities_model_BoW_cwe_capec.pkl\"\n",
    "with open(filename, \"rb\") as f:\n",
    "    BoW_clf = pickle.load(f)\n",
    "\n",
    "filename = \"link_probabilities_model_spaCy_cwe_capec.pkl\"\n",
    "with open(filename, \"rb\") as f:\n",
    "    spaCy_clf = pickle.load(f)\n",
    "\n",
    "\"\"\"filename = 'link_probabilities_model_BERT.pkl'\n",
    "with open(filename, \"rb\") as f:\n",
    "    BERT_clf = pickle.load(f)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = []\n",
    "\n",
    "for ap in ap_dict:\n",
    "    for cwe in w_dict:\n",
    "        examples.append((ap, cwe))\n",
    "\n",
    "examples = list(set(examples) - set(example_ids))\n",
    "BoW_negative_input = []\n",
    "spaCy_negative_input = []\n",
    "BERT_negative_input = []\n",
    "\n",
    "all_relevant_text = []\n",
    "for (capec, cwe) in examples:\n",
    "    spaCy_negative_input.append(handle_data(\"spaCy\", \"A0\", capec, cwe))\n",
    "    # BERT_negative_input.append(handle_data(\"BERT\", \"A1\", capec, technique, bert_output_type=\"hidden_state\", bert_finetuned=True))\n",
    "    all_relevant_text.append(gather_text_data(capec, cwe, \"A0\"))\n",
    "\n",
    "BoW_y_probs = BoW_clf.predict_proba(BoW_negative_input)[:, 1]\n",
    "spaCy_y_probs = spaCy_clf.predict_proba(spaCy_negative_input)[:, 1]\n",
    "# BERT_y_probs = BERT_clf.predict_proba(BERT_negative_input)[:, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe_dict = {}\n",
    "for i, (capec, cwe) in enumerate(examples):\n",
    "    dataframe_dict[(capec, cwe)] = {\n",
    "        \"CAPEC\": capec.split(\"/\")[1],\n",
    "        \"CWE\": cwe.split(\"/\")[1],\n",
    "        \"CAPEC name\": ap_dict[capec][\"name\"],\n",
    "        \"CWE name\": w_dict[cwe][\"name\"],\n",
    "        \"BoW link probability\": BoW_y_probs[i],\n",
    "        \"spaCy link probability\": spaCy_y_probs[i],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(dataframe_dict).T\n",
    "df.reset_index(drop=True, inplace=True)\n",
    "df.to_csv(\"link_probs_spaCy_cwe_capec.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "fe5398c8201012d3ec3837fc46b21f6bb5cc60fe8bdaa4e264ad4d58fc1758ab"
  },
  "kernelspec": {
   "display_name": "Python 3.6.9 ('venv_BRON_ML')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}