MITREembed_spotlight.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c59d4f4d-eb4e-4665-aa36-4e081154bb3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4c5f023c-104a-438f-9093-620d6fe54882",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-30 20:08:37.613727: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:08:37.616105: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:08:37.617874: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:08:37.619692: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:08:37.621440: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:08:37.623150: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1919] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13719 MB memory:  -> device: 0, name: Tesla T4, pci bus id: 0000:00:1e.0, compute capability: 7.5\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "config = tf.compat.v1.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "session = tf.compat.v1.Session(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d7d1696b-720e-4be4-b700-a1d98cdec3ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "sentences = [\"This is an example sentence\", \"Each sentence is converted\"]\n",
    "\n",
    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')\n",
    "embeddings = model.encode(sentences)\n",
    "#print(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b69b196-5445-4a4e-aede-a361966bfe4f",
   "metadata": {},
   "source": [
    "## load mitre sigma master csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "48186f9e-8306-4dac-a8ac-424f6761f87f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "mitre_sigma = pd.read_csv(\"MITRE/mitreembed_master_Chroma.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ad41c2ba-7c9c-4a91-84ca-dd55ecb8d874",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Source\n",
       "SIGMA      3047\n",
       "MITRE       566\n",
       "Netflow      13\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mitre_sigma.Source.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "991295ad-7071-446c-a9f1-a966aea93e4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "# download embeddings model\n",
    "original_model = SentenceTransformer('all-MiniLM-L12-v2')\n",
    "# reload model using langchain wrapper\n",
    "original_model.save('./')\n",
    "\n",
    "embedding_model_path = './'\n",
    "embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "dc1fa4bb-d6ec-42b1-af40-1a625b5f18fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#load model\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L12-v2')\n",
    "model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L12-v2')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3129a5d6-c08f-4b70-8701-32c503cad33f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f89c3500-e065-49df-befb-dc6a9b6ee3c8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.4.10'"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#set chromadb \n",
    "#import pydantic\n",
    "#from pydantic_settings import BaseSettings\n",
    "\n",
    "from langchain.vectorstores import Chroma\n",
    "# define logic for embeddings storage\n",
    "chromadb_path = './'\n",
    "import chromadb\n",
    "chroma_client = chromadb.Client()\n",
    "chroma_client.get_version()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a7d15d09-9df0-4b45-be28-f615661d34ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    " \n",
    "from langchain.document_loaders import DataFrameLoader\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "import re\n",
    "#import mlflow\n",
    " \n",
    "import pandas as pd\n",
    "# assemble product documents in required format (id, text)\n",
    "documents = (\n",
    "  DataFrameLoader(\n",
    "    #CISA_df_pd,\n",
    "    mitre_sigma,\n",
    "    page_content_column='Body'\n",
    "    )\n",
    "    .load()\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e73a1cd2-4807-4e55-9f2b-4b9996f0fd4b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define logic for embeddings storage\n",
    "chromadb_path = './'\n",
    "\n",
    "vectordb = Chroma.from_documents(\n",
    "  documents=documents, \n",
    "  embedding=embedding_model, \n",
    "  persist_directory=chromadb_path, \n",
    "  #collection_name = 'CISA_MITRE'\n",
    "  )\n",
    " \n",
    "# persist vector db to storage\n",
    "vectordb.persist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4e6e527-3cfd-43ca-a315-b6f01e1450b1",
   "metadata": {},
   "source": [
    "## check the number of collections in the vectordb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3ecb1449-ffae-45c3-8ef7-131826f4a249",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3626"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#delete vectors\n",
    "#vectordb.delete_collection()\n",
    "\n",
    "#count documents \n",
    "vectordb._collection.count()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10a64382-2037-43b6-9eda-aa3dc3893a7f",
   "metadata": {},
   "source": [
    "## Load CISA Advisories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0ecdf615-1539-44a2-a5d1-946d7cde6aec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "df_CISA = pd.read_csv(\"CISA_combo_features_new_feb.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "098a9d7a-366a-4bb4-adbe-1f66cbd76e97",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#add source field \n",
    "df_CISA['Source'] = 'CISA Advisory'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b91a2bcb-4d6e-4ebf-bb30-9dd3c848a722",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Subject</th>\n",
       "      <th>Date</th>\n",
       "      <th>Body</th>\n",
       "      <th>filepath</th>\n",
       "      <th>Source</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>Tue, 30 Jan 2024 11:51:59 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>Mon, 22 Jan 2024 21:43:22 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>Thu, 18 Jan 2024 12:15:57 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of January ...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>58</th>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>Wed, 03 Jan 2024 14:27:32 +0000</td>\n",
       "      <td>Vulnerability Summary for the Week of Dece...</td>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>60</th>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>Tue, 26 Dec 2023 23:25:17 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>68</th>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>Tue, 19 Dec 2023 13:17:49 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of December...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>130</th>\n",
       "      <td>Vulnerability Summary for the Week of October ...</td>\n",
       "      <td>Mon, 06 Nov 2023 22:48:34 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of October ...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>147</th>\n",
       "      <td>Vulnerability Summary for the Week of October ...</td>\n",
       "      <td>Tue, 24 Oct 2023 16:05:45 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of October ...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>181</th>\n",
       "      <td>Vulnerability Summary for the Week of Septembe...</td>\n",
       "      <td>Tue, 03 Oct 2023 12:26:33 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of Septembe...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>211</th>\n",
       "      <td>Vulnerability Summary for the Week of Septembe...</td>\n",
       "      <td>Mon, 11 Sep 2023 20:50:46 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of Septembe...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>240</th>\n",
       "      <td>Vulnerability Summary for the Week of August 7...</td>\n",
       "      <td>Mon, 14 Aug 2023 21:57:12 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of August 7...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>268</th>\n",
       "      <td>Vulnerability Summary for the Week of July 10,...</td>\n",
       "      <td>Wed, 19 Jul 2023 14:45:59 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of July 10,...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>292</th>\n",
       "      <td>Vulnerability Summary for the Week of June 26,...</td>\n",
       "      <td>Mon, 03 Jul 2023 20:40:47 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of June 26,...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>305</th>\n",
       "      <td>Vulnerability Summary for the Week of June 12,...</td>\n",
       "      <td>Thu, 22 Jun 2023 13:31:19 +0000</td>\n",
       "      <td>Cybersecurity and Infrastructure Security Agen...</td>\n",
       "      <td>Vulnerability Summary for the Week of June 12,...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               Subject  \\\n",
       "17   Vulnerability Summary for the Week of January ...   \n",
       "28   Vulnerability Summary for the Week of January ...   \n",
       "38   Vulnerability Summary for the Week of January ...   \n",
       "58   Vulnerability Summary for the Week of December...   \n",
       "60   Vulnerability Summary for the Week of December...   \n",
       "68   Vulnerability Summary for the Week of December...   \n",
       "130  Vulnerability Summary for the Week of October ...   \n",
       "147  Vulnerability Summary for the Week of October ...   \n",
       "181  Vulnerability Summary for the Week of Septembe...   \n",
       "211  Vulnerability Summary for the Week of Septembe...   \n",
       "240  Vulnerability Summary for the Week of August 7...   \n",
       "268  Vulnerability Summary for the Week of July 10,...   \n",
       "292  Vulnerability Summary for the Week of June 26,...   \n",
       "305  Vulnerability Summary for the Week of June 12,...   \n",
       "\n",
       "                                Date  \\\n",
       "17   Tue, 30 Jan 2024 11:51:59 +0000   \n",
       "28   Mon, 22 Jan 2024 21:43:22 +0000   \n",
       "38   Thu, 18 Jan 2024 12:15:57 +0000   \n",
       "58   Wed, 03 Jan 2024 14:27:32 +0000   \n",
       "60   Tue, 26 Dec 2023 23:25:17 +0000   \n",
       "68   Tue, 19 Dec 2023 13:17:49 +0000   \n",
       "130  Mon, 06 Nov 2023 22:48:34 +0000   \n",
       "147  Tue, 24 Oct 2023 16:05:45 +0000   \n",
       "181  Tue, 03 Oct 2023 12:26:33 +0000   \n",
       "211  Mon, 11 Sep 2023 20:50:46 +0000   \n",
       "240  Mon, 14 Aug 2023 21:57:12 +0000   \n",
       "268  Wed, 19 Jul 2023 14:45:59 +0000   \n",
       "292  Mon, 03 Jul 2023 20:40:47 +0000   \n",
       "305  Thu, 22 Jun 2023 13:31:19 +0000   \n",
       "\n",
       "                                                  Body  \\\n",
       "17   Cybersecurity and Infrastructure Security Agen...   \n",
       "28   Cybersecurity and Infrastructure Security Agen...   \n",
       "38   Cybersecurity and Infrastructure Security Agen...   \n",
       "58       Vulnerability Summary for the Week of Dece...   \n",
       "60   Cybersecurity and Infrastructure Security Agen...   \n",
       "68   Cybersecurity and Infrastructure Security Agen...   \n",
       "130  Cybersecurity and Infrastructure Security Agen...   \n",
       "147  Cybersecurity and Infrastructure Security Agen...   \n",
       "181  Cybersecurity and Infrastructure Security Agen...   \n",
       "211  Cybersecurity and Infrastructure Security Agen...   \n",
       "240  Cybersecurity and Infrastructure Security Agen...   \n",
       "268  Cybersecurity and Infrastructure Security Agen...   \n",
       "292  Cybersecurity and Infrastructure Security Agen...   \n",
       "305  Cybersecurity and Infrastructure Security Agen...   \n",
       "\n",
       "                                              filepath         Source  \n",
       "17   Vulnerability Summary for the Week of January ...  CISA Advisory  \n",
       "28   Vulnerability Summary for the Week of January ...  CISA Advisory  \n",
       "38   Vulnerability Summary for the Week of January ...  CISA Advisory  \n",
       "58   Vulnerability Summary for the Week of December...  CISA Advisory  \n",
       "60   Vulnerability Summary for the Week of December...  CISA Advisory  \n",
       "68   Vulnerability Summary for the Week of December...  CISA Advisory  \n",
       "130  Vulnerability Summary for the Week of October ...  CISA Advisory  \n",
       "147  Vulnerability Summary for the Week of October ...  CISA Advisory  \n",
       "181  Vulnerability Summary for the Week of Septembe...  CISA Advisory  \n",
       "211  Vulnerability Summary for the Week of Septembe...  CISA Advisory  \n",
       "240  Vulnerability Summary for the Week of August 7...  CISA Advisory  \n",
       "268  Vulnerability Summary for the Week of July 10,...  CISA Advisory  \n",
       "292  Vulnerability Summary for the Week of June 26,...  CISA Advisory  \n",
       "305  Vulnerability Summary for the Week of June 12,...  CISA Advisory  "
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#df_CISA[df_CISA['Body'].str.contains(\"MITRE\")]\n",
    "df_CISA.dropna()[df_CISA.dropna()['Body'].str.contains('mitre')]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8ee6b583-1699-4489-bd59-e08040c774b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    " \n",
    "from langchain.document_loaders import DataFrameLoader\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "import re\n",
    "#import mlflow\n",
    " \n",
    "import pandas as pd\n",
    "# assemble product documents in required format (id, text)\n",
    "documents = (\n",
    "  DataFrameLoader(\n",
    "    #CISA_df_pd,\n",
    "    df_CISA,\n",
    "    page_content_column='Body'\n",
    "    )\n",
    "    .load()\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "11e25b75-c589-4ba1-8807-ce44e75086c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define logic for embeddings storage\n",
    "chromadb_path = './'\n",
    "\n",
    "vectordb = Chroma.from_documents(\n",
    "  documents=documents, \n",
    "  embedding=embedding_model, \n",
    "  persist_directory=chromadb_path, \n",
    "  #collection_name = 'CISA_MITRE'\n",
    "  )\n",
    " \n",
    "# persist vector db to storage\n",
    "vectordb.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f7221980-3c48-4463-a77f-85ec0851658b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7771"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "#delete vectors\n",
    "#vectordb.delete_collection()\n",
    "\n",
    "#count documents \n",
    "vectordb._collection.count()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0dc1282f-ecc1-4400-86a7-93321d83f545",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Subject</th>\n",
       "      <th>Date</th>\n",
       "      <th>Body</th>\n",
       "      <th>filepath</th>\n",
       "      <th>Source</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CISA Partners With OpenSSF Securing Software R...</td>\n",
       "      <td>Thu, 08 Feb 2024 19:16:35 +0000</td>\n",
       "      <td>CISA Partners With OpenSSF Securing Softwa...</td>\n",
       "      <td>CISA Partners With OpenSSF Securing Software R...</td>\n",
       "      <td>CISA Advisory</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             Subject  \\\n",
       "0  CISA Partners With OpenSSF Securing Software R...   \n",
       "\n",
       "                              Date  \\\n",
       "0  Thu, 08 Feb 2024 19:16:35 +0000   \n",
       "\n",
       "                                                Body  \\\n",
       "0      CISA Partners With OpenSSF Securing Softwa...   \n",
       "\n",
       "                                            filepath         Source  \n",
       "0  CISA Partners With OpenSSF Securing Software R...  CISA Advisory  "
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_CISA.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2218ef8f-5498-4cad-86d2-1318325301d7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metadatas:   [{'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Analyze network data for uncommon data flows (e.g., a client sending significantly more data than it receives from a server). Processes utilizing the network that do not normally have network communication or have never been seen before are suspicious. Analyze packet contents to detect communications that do not follow the expected protocol behavior for the port that is being used. (Citation: University of Birmingham C2)', 'filepath': 'https://attack.mitre.org/techniques/T1001'}]\n",
      "Documents:   ['Adversaries may obfuscate command and control traffic to make it more difficult to detect. Command and control (C2) communications are hidden (but not necessarily encrypted) in an attempt to make the content more difficult to discover or decipher and to make the communication less conspicuous and hide commands from being seen. This encompasses many methods, such as adding junk data to protocol traffic, using steganography, or impersonating legitimate protocols. ']\n",
      "ids:         ['7946afd4-072d-11ef-9f94-d742a8cd9055']\n",
      "embeddings:  [[-0.03199813514947891, 0.03938308730721474, -0.03308778256177902, -0.08212287724018097, 0.021442661061882973, -0.005276898853480816, 0.034147098660469055, -0.01586325280368328, 0.054854944348335266, 0.032911185175180435, 0.07466107606887817, 0.058753374963998795, 0.000574931560549885, -0.08866751939058304, 0.017476102337241173, 0.05606706812977791, 0.04520796611905098, -0.0209503136575222, 0.027694301679730415, -0.07381413131952286, -0.014573896303772926, -0.02202725037932396, 0.05289576202630997, 0.0035967794246971607, -0.025043269619345665, 0.021317249163985252, -0.059706974774599075, -0.05719034746289253, -7.230418850667775e-05, -0.08419245481491089, -0.01109999231994152, -0.014616911299526691, 0.07802678644657135, -0.06867887079715729, 0.05790197476744652, 0.01662546768784523, 0.058324363082647324, -0.05282101407647133, 0.036543089896440506, 0.0013748093042522669, -0.0667872205376625, -0.01398713793605566, -0.026739157736301422, 0.050916701555252075, -0.018132872879505157, -0.01641785353422165, -0.07843563705682755, 0.03396415710449219, -0.023769157007336617, -0.0038760299794375896, 0.06732284277677536, 0.0294589065015316, -0.01320599764585495, 0.01715032197535038, 0.05065358430147171, -0.07113341242074966, -0.06299860030412674, 0.12201543152332306, -0.05049475282430649, 0.0992717444896698, -0.03611095994710922, 0.010513024404644966, -0.015329229645431042, 0.02791936695575714, -0.007066342048346996, 0.06847228854894638, -0.05553826317191124, 0.021903278306126595, 0.022279448807239532, 0.06056308373808861, -0.02107083797454834, -0.00015017321857158095, -0.11538343876600266, -0.010793698951601982, 0.02615639939904213, 0.04252876713871956, -0.12814706563949585, -0.04085814207792282, 0.0014883950352668762, -0.05519422888755798, 0.005219947546720505, 0.034477151930332184, 0.012133481912314892, 0.05412069335579872, -0.030908962711691856, -0.03309090808033943, -0.013449340127408504, -0.0637541264295578, 0.07632585614919662, -0.028028268367052078, -0.05249791592359543, -0.05601835250854492, -0.030880803242325783, -0.03279316425323486, -0.06265296041965485, 0.0018446987960487604, 0.04914597049355507, -0.08142843097448349, 0.03144768252968788, 0.10264983773231506, 0.03679981827735901, -0.03744816407561302, 0.03338282182812691, -0.021922621876001358, -0.02682073600590229, 0.025262514129281044, 0.050265610218048096, -0.03952324390411377, -0.011342189274728298, -0.01102761086076498, -0.10203582048416138, 0.031771548092365265, -0.06824661791324615, -0.07430416345596313, -0.023427698761224747, -0.0174636859446764, -0.07513779401779175, 0.03514669090509415, -0.010440980084240437, 0.008153341710567474, -0.07522699236869812, -0.04336153715848923, 0.08551321923732758, -0.007389440201222897, 0.052948515862226486, -0.047749124467372894, -0.012780806981027126, 0.005923858378082514, 0.01834789477288723, -0.03133144974708557, -0.046249933540821075, -0.044689424335956573, -0.013681734912097454, 0.0324280746281147, 0.05290713906288147, -0.08693419396877289, -0.04595388472080231, 0.05500352010130882, -0.004749923478811979, -0.055145129561424255, 0.1017729789018631, 0.06540383398532867, 0.09810163825750351, -0.019443484023213387, 0.04681668430566788, 0.022299718111753464, -0.0230714064091444, -0.06904297322034836, -0.012655897997319698, -0.017534475773572922, -0.006583860144019127, 0.0020548724569380283, 0.10218594968318939, 0.021870633587241173, -0.021630706265568733, -0.01925419084727764, 0.12405651062726974, 0.029817990958690643, 0.009962527081370354, 0.05370049551129341, 0.0887569710612297, 0.0825926661491394, 0.05720853433012962, 0.07303525507450104, -0.0011020450619980693, -0.08406753838062286, 0.009431147016584873, -0.014389467425644398, -0.013091560453176498, -0.04764476791024208, -0.036733005195856094, -0.0731034129858017, -0.045632071793079376, -0.07225470244884491, -0.0031545963138341904, -0.04518596827983856, -0.10789363086223602, 0.02716229110956192, 0.015563336201012135, -0.009324191138148308, 0.011088312603533268, -0.10043071210384369, 0.06630616635084152, -0.018814710900187492, 0.021945813670754433, 0.005703335627913475, 0.017089281231164932, -0.036203332245349884, -0.00692348787561059, -0.00958920456469059, 0.03390221297740936, 0.0461861789226532, 0.024643775075674057, 0.011349847540259361, -0.04933195188641548, -0.07918567955493927, 0.05413800850510597, -0.029784685000777245, -0.11101380735635757, 0.056027449667453766, -0.035229604691267014, -0.004220633767545223, 0.021991228684782982, 0.07613925635814667, -0.09007832407951355, -0.02418588660657406, 0.0021588706877082586, -0.018004905432462692, -0.01151470746845007, 0.01812923327088356, 0.038432665169239044, 0.07727715373039246, -0.07142636924982071, -0.10547498613595963, -0.11657779663801193, -0.03240935504436493, 0.024711372330784798, 0.07621258497238159, 0.011168211698532104, 0.004118578042834997, 0.01844431459903717, 0.014173438772559166, 0.028631186112761497, 5.8265129719773375e-33, -0.030394263565540314, 0.06923531740903854, -0.011013036593794823, -0.005508314818143845, -0.10994568467140198, 0.0046081747859716415, 0.020318280905485153, 0.025725821033120155, 0.024991760030388832, -0.014910051599144936, -0.03241705149412155, 0.034893982112407684, -0.03620079532265663, -0.024400237947702408, 0.023182962089776993, -0.009270635433495045, 0.024022823199629784, 0.003432880388572812, -0.043867163360118866, 0.019609514623880386, 0.034881725907325745, 0.017919812351465225, -0.003184063360095024, -0.016709042713046074, -0.005910683888942003, -0.014472865499556065, -0.06459103524684906, 0.052251044660806656, 0.05989221855998039, 0.02516964264214039, 0.016075359657406807, 0.07557566463947296, -0.06088500842452049, 0.03023795783519745, -0.02498323656618595, -0.03161554038524628, 0.0052047292701900005, 0.07829509675502777, -0.05130380392074585, -0.08609017729759216, -0.05110599845647812, 0.003961300943046808, -0.10449930280447006, 0.10073336958885193, 0.042726099491119385, 0.002933512208983302, 0.0009666797704994678, 0.030922675505280495, -0.056453533470630646, -0.0455387718975544, 0.004968469496816397, -0.039271555840969086, 0.10137143731117249, -0.04977641627192497, -0.00021162121265660971, 0.04294627532362938, -0.09324544668197632, -0.03597675636410713, 0.06691111624240875, 0.00018344400450587273, 0.12539631128311157, -0.04369235038757324, -0.10706166177988052, 0.046285517513751984, 0.04356107860803604, 0.06444232165813446, -0.04926861450076103, 0.0732637420296669, 0.057507213205099106, -0.07169850915670395, 0.08178295195102692, -0.035594258457422256, 0.031270548701286316, -0.046613071113824844, 0.012587817385792732, -0.042263031005859375, -0.019170328974723816, -0.03487621247768402, -0.08780398964881897, 0.06336010992527008, 0.048869505524635315, 0.012639702297747135, 0.006712328642606735, -0.024500606581568718, 0.047868456691503525, -0.007346172351390123, 0.014921399764716625, 0.021391572430729866, 0.0066968342289328575, -0.040818456560373306, 0.003284849924966693, 0.010878357104957104, -0.10619471222162247, 0.04480284824967384, -0.00872271042317152, 6.390133640116632e-32, 0.07674793154001236, -0.001823689672164619, 0.03125562146306038, -0.07621476799249649, 0.04134536162018776, 0.10212186723947525, 0.005636087153106928, -0.14537478983402252, 0.00898090098053217, -0.009015975520014763, 0.049695245921611786, -0.03332745283842087, -0.014611955732107162, -0.018458178266882896, 0.024889547377824783, 0.07729648798704147, -0.10888421535491943, -0.0678003579378128, 0.05537371709942818, -0.06065112352371216, -0.047029223293066025, 0.05034404247999191, -0.05241651460528374, 0.022969646379351616, 0.005690648220479488, 0.030191779136657715, 0.010845350101590157, 0.11406387388706207, 0.030631287023425102, 0.01955675333738327, -0.029326345771551132, -0.032372504472732544, -0.037002939730882645, 0.06703409552574158, -0.03382081910967827, 0.1573258340358734, 0.04657512530684471, 0.00044340509339235723, -0.06483186036348343, 0.03902557119727135, 0.0006364720757119358, -0.0555991530418396, -0.07520990818738937, 0.04354533925652504, -0.05833888426423073, -0.04781443253159523, 0.04143421724438667, 0.012894764542579651, 0.03779294714331627, 0.019265245646238327, -0.022972064092755318, 0.03263038024306297, -0.023806996643543243, 0.08161727339029312, 0.05519217997789383, -0.004629986826330423, 0.03644152730703354, -0.03440001606941223, 0.03908780962228775, 0.059123553335666656, -0.07030925154685974, 0.06165247783064842, 0.025969045236706734, 0.014046313241124153]]\n"
     ]
    }
   ],
   "source": [
    "#examine a vector db record\n",
    "rec= vectordb._collection.peek(1)\n",
    "print('Metadatas:  ', rec['metadatas'])\n",
    "print('Documents:  ', rec['documents'])\n",
    "print('ids:        ', rec['ids']) \n",
    "print('embeddings: ', rec['embeddings'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "26fd332c-5a77-46d5-9c86-a5862b555500",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#results = vectordb._collection.query(include=['embeddings'])\n",
    "#results_json = vectordb.get()\n",
    "results_json = vectordb.get(include=['embeddings', 'documents', 'metadatas'])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "39982948-9868-4fbe-9218-4b605c43e569",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "#data = json.loads(results_json).astype(str)\n",
    "#embeddings = json.dumps(data)\n",
    "#pd.read_json(embeddings)\n",
    "embeddings = pd.DataFrame.from_dict(results_json)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "10850bb7-7713-4c51-9896-419219ed0cd4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "embeddings.to_csv('mitre_embeddings.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8ed1b660-6c10-41ae-ae57-b0a84d99401e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(Document(page_content=\"Adversaries may move onto systems, possibly those on disconnected or air-gapped networks, by copying malware to removable media and taking advantage of Autorun features when the media is inserted into a system and executes. In the case of Lateral Movement, this may occur through modification of executable files stored on removable media or by copying malware and renaming it to look like a legitimate file to trick users into executing it on a separate system. In the case of Initial Access, this may occur through manual manipulation of the media, modification of systems used to initially format the media, or modification to the media's firmware itself.\", metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Monitor file access on removable media. Detect processes that execute from removable media after it is mounted or when initiated by a user. If a remote access tool is used in this manner to move laterally, then additional actions are likely to occur after execution, such as opening network connections for Command and Control and system and network information Discovery.', 'filepath': 'https://attack.mitre.org/techniques/T1091'}),\n",
       "  1.171260118484497),\n",
       " (Document(page_content='Adversaries may transfer tools or other files from an external system into a compromised environment. Files may be copied from an external adversary controlled system through the command and control channel to bring tools into the victim network or through alternate protocols with another tool such as FTP. Files can also be copied over on Mac and Linux with native tools like scp, rsync, and sftp.', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Monitor for file creation and files transferred into the network. Unusual processes with external network connections creating files on-system may be suspicious. Use of utilities, such as FTP, that does not normally occur may also be suspicious.\\n\\nAnalyze network data for uncommon data flows (e.g., a client sending significantly more data than it receives from a server). Processes utilizing the network that do not normally have network communication or have never been seen before are suspicious. Analyze packet contents to detect communications that do not follow the expected protocol behavior for the port that is being used.(Citation: University of Birmingham C2)', 'filepath': 'https://attack.mitre.org/techniques/T1105'}),\n",
       "  1.2550523281097412),\n",
       " (Document(page_content='Adversaries may transfer tools or other files between systems in a compromised environment. Files may be copied from one system to another to stage adversary tools or other files over the course of an operation. Adversaries may copy files laterally between internal victim systems to support lateral movement using inherent file sharing protocols such as file sharing over SMB to connected network shares or with authenticated connections with [SMB/Windows Admin Shares](T1021.002) or [Remote Desktop Protocol](T1021.001). Files can also be copied over on Mac and Linux with native tools like scp, rsync, and sftp.', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Monitor for file creation and files transferred within a network using protocols such as SMB. Unusual processes with internal network connections creating files on-system may be suspicious. Consider monitoring for abnormal usage of utilities and command-line arguments that may be used in support of remote transfer of files. Considering monitoring for alike file hashes or characteristics (ex: filename) that are created on multiple hosts.', 'filepath': 'https://attack.mitre.org/techniques/T1570'}),\n",
       "  1.2785615921020508),\n",
       " (Document(page_content=\"\\nAdversaries may deliver payloads to remote systems by adding content to shared storage locations, such as network drives or internal code repositories. Content stored on network drives or in other shared locations may be tainted by adding malicious programs, scripts, or exploit code to otherwise valid files. Once a user opens the shared tainted content, the malicious portion can be executed to run the adversary's code on a remote system. Adversaries may use tainted shared content to move laterally.\\n\\nA directory share pivot is a variation on this technique that uses several other techniques to propagate malware when users access a shared network directory. It uses [Shortcut Modification](T1547.009) of directory .LNK files that use [Masquerading](T1036) to look like the real directories, which are hidden through [Hidden Files and Directories](T1564.001). The malicious .LNK-based directories have an embedded command that executes the hidden malware file in the directory and then opens the real intended directory so that the user's expected action still occurs. When used with frequently used network directories, the technique may result in frequent reinfections and broad access to systems and potentially to new and higher privileged accounts. (Citation: Retwin Directory Share Pivot)\\n\\nAdversaries may also compromise shared network directories through binary infections by appending or prepending its code to the healthy binary on the shared network directory. The malware may modify the original entry point (OEP) of the healthy binary to ensure that it is executed before the legitimate code. The infection could continue to spread via the newly infected file when it is executed by a remote system. These infections may target both binary and non-binary formats that end with extensions including, but not limited to, .EXE, .DLL, .SCR, .BAT, and/or .VBS.\", metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Processes that write or overwrite many files to a network shared directory may be suspicious. Monitor processes that are executed from removable media for malicious or abnormal activity such as network connections due to Command and Control and possible network Discovery techniques.\\n\\nFrequently scan shared network directories for malicious files, hidden files, .LNK files, and other file types that may not typical exist in directories used to share specific types of content.', 'filepath': 'https://attack.mitre.org/techniques/T1080'}),\n",
       "  1.3220890760421753)]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vectordb.similarity_search_with_score(\"MoveIt Vulnerability\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "cbf7228c-49e5-426f-8218-d3d0088d52d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "##vectordb.similarity_search_with_relevance_scores(\"MoveIt Vulnerability\", k=50, score_threshold=0.80, search_type=\"hybrid\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "a7547307-f772-4bc1-b9f1-ecdaa9854f20",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ProductSearchWrapper:\n",
    "    def __init__(self, embedding_model, chromadb, max_results=25):\n",
    "        # Retrieve embedding model\n",
    "        self.embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_path)\n",
    "\n",
    "        # Retrieve vectordb contents\n",
    "        self._vectordb = Chroma(persist_directory=chromadb_path, embedding_function=self.embedding_model)\n",
    "\n",
    "        # Set number of results to returnd\n",
    "        self._max_results = max_results\n",
    "\n",
    "    def predict(self, query):\n",
    "        # Perform search on embeddings\n",
    "        raw_results = self._vectordb.similarity_search_with_score(query, k=self._max_results)\n",
    "\n",
    " \n",
    "        # get lists of of scores, descriptions and ids from raw results\n",
    "        scores, Body, Subject, Date, filepath, Source = zip(\n",
    "          *[(r[1], r[0].page_content, r[0].metadata['Subject'], r[0].metadata['Date'], r[0].metadata['filepath'], \n",
    "             r[0].metadata['Source']) \n",
    "             for r in raw_results] \n",
    "          )\n",
    "        results_pd = pd.DataFrame({        \n",
    "            #'Body':Body,\n",
    "            'Subject':Subject,\n",
    "            'Date':Date,\n",
    "            'Body':Body,\n",
    "            'filepath':filepath,\n",
    "            'Source':Source,\n",
    "            'score':scores\n",
    "          }).sort_values(axis=0, by='score', ascending=True)\n",
    "\n",
    "        #set return value\n",
    "        return results_pd\n",
    "        #return raw_results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "a4aceb55-aea3-44e6-b6a5-956e75c85fe6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Subject</th>\n",
       "      <th>Date</th>\n",
       "      <th>Body</th>\n",
       "      <th>filepath</th>\n",
       "      <th>Source</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Monitor file access on removable media. Detect...</td>\n",
       "      <td>4/5/24</td>\n",
       "      <td>Adversaries may move onto systems, possibly th...</td>\n",
       "      <td>https://attack.mitre.org/techniques/T1091</td>\n",
       "      <td>MITRE</td>\n",
       "      <td>1.171260</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Monitor for file creation and files transferre...</td>\n",
       "      <td>4/5/24</td>\n",
       "      <td>Adversaries may transfer tools or other files ...</td>\n",
       "      <td>https://attack.mitre.org/techniques/T1105</td>\n",
       "      <td>MITRE</td>\n",
       "      <td>1.255052</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Monitor for file creation and files transferre...</td>\n",
       "      <td>4/5/24</td>\n",
       "      <td>Adversaries may transfer tools or other files ...</td>\n",
       "      <td>https://attack.mitre.org/techniques/T1570</td>\n",
       "      <td>MITRE</td>\n",
       "      <td>1.278562</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Processes that write or overwrite many files t...</td>\n",
       "      <td>4/5/24</td>\n",
       "      <td>\\nAdversaries may deliver payloads to remote s...</td>\n",
       "      <td>https://attack.mitre.org/techniques/T1080</td>\n",
       "      <td>MITRE</td>\n",
       "      <td>1.322089</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Monitor for common command-line editors used t...</td>\n",
       "      <td>4/5/24</td>\n",
       "      <td>Adversaries can modify property list files (pl...</td>\n",
       "      <td>https://attack.mitre.org/techniques/T1547/011</td>\n",
       "      <td>MITRE</td>\n",
       "      <td>1.401985</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             Subject    Date  \\\n",
       "0  Monitor file access on removable media. Detect...  4/5/24   \n",
       "1  Monitor for file creation and files transferre...  4/5/24   \n",
       "2  Monitor for file creation and files transferre...  4/5/24   \n",
       "3  Processes that write or overwrite many files t...  4/5/24   \n",
       "4  Monitor for common command-line editors used t...  4/5/24   \n",
       "\n",
       "                                                Body  \\\n",
       "0  Adversaries may move onto systems, possibly th...   \n",
       "1  Adversaries may transfer tools or other files ...   \n",
       "2  Adversaries may transfer tools or other files ...   \n",
       "3  \\nAdversaries may deliver payloads to remote s...   \n",
       "4  Adversaries can modify property list files (pl...   \n",
       "\n",
       "                                        filepath Source     score  \n",
       "0      https://attack.mitre.org/techniques/T1091  MITRE  1.171260  \n",
       "1      https://attack.mitre.org/techniques/T1105  MITRE  1.255052  \n",
       "2      https://attack.mitre.org/techniques/T1570  MITRE  1.278562  \n",
       "3      https://attack.mitre.org/techniques/T1080  MITRE  1.322089  \n",
       "4  https://attack.mitre.org/techniques/T1547/011  MITRE  1.401985  "
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Instantiate the model\n",
    "#model objective:\n",
    "\"\"\"\n",
    "The model is intented to be used as a sentence and short paragraph encoder. \n",
    "Given an input text, it ouptuts a vector which captures the semantic information. \n",
    "The sentence vector may be used for information retrieval, clustering or sentence similarity tasks.\n",
    "\"\"\"\n",
    "\n",
    "model = ProductSearchWrapper(embedding_model=embedding_model_path, chromadb=chromadb_path, max_results=25)\n",
    "\n",
    "\n",
    "# Call the model's predict method with a query\n",
    "#query = \"Which is the most recent CISA advisory on Ivanti\"\n",
    "query = \"MoveIT vulnerability\"\n",
    "results = model.predict(query)\n",
    "\n",
    "#top 5 results\n",
    "df_results = pd.DataFrame(results).head(5)\n",
    "#df_results\n",
    "\n",
    "#df_results.to_csv('result_test.csv')\n",
    "# response_text = df_results[['Subject', 'Date','filepath','score']]\n",
    "# response_text\n",
    "#str(response_text.iloc[0]['Subject'])\n",
    "\n",
    "#df_results.iloc[0]\n",
    "#df_results.iloc[0]['Body']\n",
    "df_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c84cf75b-e75f-4f04-9609-06d1b18f23e0",
   "metadata": {},
   "source": [
    "## LLM integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "2c97039b-5e86-4b66-82bb-fd91558f1bfb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__CUDNN VERSION: 8904\n",
      "__Number CUDA Devices: 1\n",
      "__CUDA Device Name: Tesla T4\n",
      "__CUDA Device Total Memory [GB]: 15.655829504\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "use_cuda = torch.cuda.is_available()\n",
    "if use_cuda:\n",
    "    print('__CUDNN VERSION:', torch.backends.cudnn.version())\n",
    "    print('__Number CUDA Devices:', torch.cuda.device_count())\n",
    "    print('__CUDA Device Name:',torch.cuda.get_device_name(0))\n",
    "    print('__CUDA Device Total Memory [GB]:',torch.cuda.get_device_properties(0).total_memory/1e9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "69dcbbcb-2482-48ed-b6bd-d806349de5fe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import torch\n",
    "\n",
    "gc.collect()\n",
    "\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "a85cf59f-840d-45fa-a29b-a8507f0313d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## initialize LLM wizardlm-13b\n",
    "from langchain.llms import LlamaCpp\n",
    "from langchain.callbacks.manager import CallbackManager\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'\n",
    "os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "9b36cde3-7159-412d-bc63-671ff0126651",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ggml_init_cublas: GGML_CUDA_FORCE_MMQ:   no\n",
      "ggml_init_cublas: CUDA_USE_TENSOR_CORES: yes\n",
      "ggml_init_cublas: found 1 CUDA devices:\n",
      "  Device 0: Tesla T4, compute capability 7.5, VMM: yes\n"
     ]
    }
   ],
   "source": [
    "n_gpu_layers = 35  # Metal set to 1 is enough.\n",
    "#n_batch = 1042  # Should be between 1 and n_ctx, consider the amount of RAM of your Apple Silicon Chip.\n",
    "callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])\n",
    "\n",
    "# Make sure the model path is correct for your system!\n",
    "llm = LlamaCpp(\n",
    "    #model_path=\"Lilly/Lily-7B-Instruct-v0.2.Q5_K_M.gguf\",\n",
    "    model_path=\"./wizardlm-13b-v1.1-superhot-8k.ggmlv3.q4_0.gguf.bin\",\n",
    "    n_gpu_layers=n_gpu_layers,\n",
    " #   n_batch=n_batch,\n",
    " #   temperature=0.75,\n",
    "    max_tokens=2000,\n",
    " #   top_p=1,\n",
    "    n_ctx=2048,\n",
    "    f16_kv=True,  # MUST set to True, otherwise you will run into problem after a couple of calls\n",
    "    callback_manager=callback_manager,\n",
    "    verbose=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "816ae02b-944d-4c93-aad5-eca9a211bf3c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/device:GPU:0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-30 20:14:39.645178: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:14:39.647554: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:14:39.649300: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:14:39.651097: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:14:39.652841: I external/local_xla/xla/stream_executor/cuda/cuda_gpu_executor.cc:899] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-04-30 20:14:39.654542: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1919] Created device /device:GPU:0 with 13719 MB memory:  -> device: 0, name: Tesla T4, pci bus id: 0000:00:1e.0, compute capability: 7.5\n"
     ]
    }
   ],
   "source": [
    "#%tensorflow_version 2.x\n",
    "#!pip install tf-nightly\n",
    "import tensorflow as tf\n",
    "device_name = tf.test.gpu_device_name()\n",
    "print(device_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d23bcf44-5af7-4407-8d29-4db753c75862",
   "metadata": {},
   "source": [
    "## intialize koboldcpp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "3c68d8b8-db3c-4f55-9d9f-0b3c844327da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#https://medium.com/p/9f890e6960f3\n",
    "# Successfully uninstalled langchain-0.0.198\n",
    "\n",
    "from langchain import PromptTemplate, LLMChain\n",
    "\n",
    "\n",
    "vectordb = Chroma(persist_directory=chromadb_path, embedding_function=embedding_model)\n",
    "#retrievers\n",
    "\n",
    "retriever = vectordb.as_retriever()\n",
    "\n",
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "# qa_chain = RetrievalQA.from_chain_type(\n",
    "#     llm,\n",
    "#     retriever=vectordb.as_retriever()\n",
    "# )\n",
    "\n",
    "# Build prompt\n",
    "template = \"\"\"Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say \"thanks for asking!\" at the end of the answer. \n",
    "{context}\n",
    "Question: {question}\n",
    "Helpful Answer:\"\"\"\n",
    "QA_CHAIN_PROMPT = PromptTemplate.from_template(template)# Run chain\n",
    "\n",
    "from langchain.chains import RetrievalQA, ConversationalRetrievalChain\n",
    "qa = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    chain_type=\"refine\",\n",
    "    #chain_type_kwargs={\"prompt\": QA_CHAIN_PROMPT},\n",
    "    retriever=vectordb.as_retriever(search_kwargs={\"k\": 1}),\n",
    "    return_source_documents=True,\n",
    "    verbose=False,\n",
    ")\n",
    "#from langchain.llms import KoboldApiLLM\n",
    "#llm = KoboldApiLLM(endpoint=\"http://localhost:5001\", max_length=80)\n",
    "#llm = KoboldApiLLM(endpoint=\"http://54.165.180.47:5001\", max_length=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "bcf179b6-d2a0-4967-985b-5f867ea95578",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------------------------------------------------------------------------------"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me what  Kerberoasting is and find a MITRE technique associated with it',\n",
       " 'result': '---------------------------------------------------------------------------------------------------------',\n",
       " 'source_documents': [Document(page_content='Adversaries may attempt to subvert Kerberos authentication by stealing or forging Kerberos tickets to enable [Pass the Ticket](T1550.003). Kerberos is an authentication protocol widely used in modern Windows domain environments. In Kerberos environments, referred to as ‚Äúrealms‚Äù, there are three basic participants: client, service, and Key Distribution Center (KDC).(Citation: ADSecurity Kerberos Ring Decoder) Clients request access to a service and through the exchange of Kerberos tickets, originating from KDC, they are granted access after having successfully authenticated. The KDC is responsible for both authentication and ticket granting.  Attackers may attempt to abuse Kerberos by stealing tickets or forging tickets to enable unauthorized access.\\n\\nOn Windows, the built-in `klist` utility can be used to list and analyze cached Kerberos tickets.(Citation: Microsoft Klist)\\n\\nLinux systems on Active Directory domains store Kerberos credentials locally in the credential cache file referred to as the \"ccache\". The credentials are stored in the ccache file while they remain valid and generally while a user\\'s session lasts.(Citation: MIT ccache) On modern Redhat Enterprise Linux systems, and derivative distributions, the System Security Services Daemon (SSSD) handles Kerberos tickets. By default SSSD maintains a copy of the ticket database that can be found in `/var/lib/sss/secrets/secrets.ldb` as well as the corresponding key located in `/var/lib/sss/secrets/.secrets.mkey`. Both files require root access to read. If an adversary is able to access the database and key, the credential cache Kerberos blob can be extracted and converted into a usable Kerberos ccache file that adversaries may use for [Pass the Ticket](T1550.003). The ccache file may also be converted into a Windows format using tools such as Kekeo.(Citation: Linux Kerberos Tickets)(Citation: Brining MimiKatz to Unix)(Citation: Kekeo)\\n\\n\\nKerberos tickets on macOS are stored in a standard ccache format, similar to Linux. By default, access to these ccache entries is federated through the KCM daemon process via the Mach RPC protocol, which uses the caller\\'s environment to determine access. The storage location for these ccache entries is influenced by the `/etc/krb5.conf` configuration file and the `KRB5CCNAME` environment variable which can specify to save them to disk or keep them protected via the KCM daemon. Users can interact with ticket storage using `kinit`, `klist`, `ktutil`, and `kcc` built-in binaries or via Apple\\'s native Kerberos framework. Adversaries can use open source tools to interact with the ccache files directly or to use the Kerberos framework to call lower-level APIs for extracting the user\\'s TGT or Service Tickets.(Citation: SpectorOps Bifrost Kerberos macOS 2019)(Citation: macOS kerberos framework MIT)\\n', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'Monitor for anomalous Kerberos activity, such as malformed or blank fields in Windows logon/logoff events (Event ID 4624, 4672, 4634), RC4 encryption within ticket granting tickets (TGTs), and ticket granting service (TGS) requests without preceding TGT requests.(Citation: ADSecurity Detecting Forged Tickets)(Citation: Stealthbits Detect PtT 2019)(Citation: CERT-EU Golden Ticket Protection)\\n\\nMonitor the lifetime of TGT tickets for values that differ from the default domain duration.(Citation: Microsoft Kerberos Golden Ticket)\\n\\nMonitor for indications of [Pass the Ticket](T1550.003) being used to move laterally. \\n\\nEnable Audit Kerberos Service Ticket Operations to log Kerberos TGS service ticket requests. Particularly investigate irregular patterns of activity (ex: accounts making numerous requests, Event ID 4769, within a small time frame, especially if they also request RC4 encryption [Type 0x17]).(Citation: Microsoft Detecting Kerberoasting Feb 2018) (Citation: AdSecurity Cracking Kerberos Dec 2015)\\n\\nMonitor for unexpected processes interacting with lsass.exe.(Citation: Medium Detecting Attempts to Steal Passwords from Memory) Common credential dumpers such as [Mimikatz](S0002) access the LSA Subsystem Service (LSASS) process by opening the process, locating the LSA secrets key, and decrypting the sections in memory where credential details, including Kerberos tickets, are stored.\\n\\nMonitor for unusual processes accessing¬†`secrets.ldb` and `.secrets.mkey` located in `/var/lib/sss/secrets/`.', 'filepath': 'https://attack.mitre.org/techniques/T1558'})]}"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me what  Kerberoasting is and find a MITRE technique associated with it\")\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "4bb2a4f0-1435-4091-90fd-f30d56eeebd9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1. The MITRE technique being used here is `remoteconnectionmanager`."
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': '  Below is a script to look through logs for malicious events, explain to me the MITRE technique for which this detection b belongs to remoteconnectionmanager` EventCode=1149 | stats count min(_time) as firstTime max(_time) as lastTime by ComputerName, Source_Network_Address, User, Message | `security_content_ctime(firstTime)` | `security_content_ctime(lastTime)` | rename ComputerName as dest | `windows_rdp_connection_successful_filter',\n",
       " 'result': '\\n1. The MITRE technique being used here is `remoteconnectionmanager`.',\n",
       " 'source_documents': [Document(page_content='Adversaries may check for Internet connectivity on compromised systems. This may be performed during automated discovery and can be accomplished in numerous ways such as using [Ping](S0097), `tracert`, and GET requests to websites.\\n\\nAdversaries may use the results and responses from these requests to determine if the system is capable of communicating with their C2 servers before attempting to connect to them. The results may also be used to identify routes, redirectors, and proxy servers.', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'System and network discovery techniques normally occur throughout an operation as an adversary learns the environment. Data and events should not be viewed in isolation, but as part of a chain of behavior that could lead to other activities, such as Command and Control, based on the information obtained.\\n\\nMonitor processes and command-line arguments for actions that could be taken to check Internet connectivity.', 'filepath': 'https://attack.mitre.org/techniques/T1016/001'})]}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"  Below is a script to look through logs for malicious events, explain to me the MITRE technique for which this detection b belongs to \\\n",
    "remoteconnectionmanager` EventCode=1149 \\\n",
    "| stats count min(_time) as firstTime max(_time) as lastTime by ComputerName, Source_Network_Address, User, Message \\\n",
    "| `security_content_ctime(firstTime)` \\\n",
    "| `security_content_ctime(lastTime)` \\\n",
    "| rename ComputerName as dest \\\n",
    "| `windows_rdp_connection_successful_filter\")\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "390f97ea-d323-48ad-aae6-3fcb6135afbc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The detection belongs to MITRE TECHNICAL OBJECT CONCEPT (TOCTOU) technique. The command is using PowerShell's `System.Net.WebClient` class along with `UploadString` method to download and execute potentially malicious code from remote locations. The script then calculates the first and last occurrence times of such events using the `security_content_ctime()` function, which could indicate continuous exfiltration of sensitive data or content over time via C2 communication channels established via PowerShell and `UploadString`."
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me Which MITRE technique does this detection belongs to powershell EventCode=4104 ScriptBlockText = *Net.webclient* AND ScriptBlockText = *.UploadString* | stats count min(_time) as firstTime max(_time) as lastTime by EventCode ScriptBlockText Computer UserID | security_content_ctime(firstTime) | security_content_ctime(lastTime) | windows_exfiltration_over_c2_via_powershell_uploadstring_filter',\n",
       " 'result': \"\\nThe detection belongs to MITRE TECHNICAL OBJECT CONCEPT (TOCTOU) technique. The command is using PowerShell's `System.Net.WebClient` class along with `UploadString` method to download and execute potentially malicious code from remote locations. The script then calculates the first and last occurrence times of such events using the `security_content_ctime()` function, which could indicate continuous exfiltration of sensitive data or content over time via C2 communication channels established via PowerShell and `UploadString`.\",\n",
       " 'source_documents': [Document(page_content=\"Adversaries may abuse PowerShell commands and scripts for execution. PowerShell is a powerful interactive command-line interface and scripting environment included in the Windows operating system. (Citation: TechNet PowerShell) Adversaries can use PowerShell to perform a number of actions, including discovery of information and execution of code. Examples include the `Start-Process` cmdlet which can be used to run an executable and the `Invoke-Command` cmdlet which runs a command locally or on a remote computer (though administrator permissions are required to use PowerShell to connect to remote systems).\\n\\nPowerShell may also be used to download and run executables from the Internet, which can be executed from disk or in memory without touching disk.\\n\\nA number of PowerShell-based offensive testing tools are available, including [Empire](S0363),  [PowerSploit](S0194), [PoshC2](S0378), and PSAttack.(Citation: Github PSAttack)\\n\\nPowerShell commands/scripts can also be executed without directly invoking the `powershell.exe` binary through interfaces to PowerShell's underlying `System.Management.Automation` assembly DLL exposed through the .NET framework and Windows Common Language Interface (CLI). (Citation: Sixdub PowerPick Jan 2016)(Citation: SilentBreak Offensive PS Dec 2015)(Citation: Microsoft PSfromCsharp APR 2014)\", metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'If proper execution policy is set, adversaries will likely be able to define their own execution policy if they obtain administrator or system access, either through the Registry or at the command line. This change in policy on a system may be a way to detect malicious use of PowerShell. If PowerShell is not used in an environment, then simply looking for PowerShell execution may detect malicious activity.\\n\\nMonitor for loading and/or execution of artifacts associated with PowerShell specific assemblies, such as System.Management.Automation.dll (especially to unusual process names/locations).(Citation: Sixdub PowerPick Jan 2016)(Citation: SilentBreak Offensive PS Dec 2015)\\n\\nIt is also beneficial to turn on PowerShell logging to gain increased fidelity in what occurs during execution (which is applied to .NET invocations). (Citation: Malware Archaeology PowerShell Cheat Sheet) PowerShell 5.0 introduced enhanced logging capabilities, and some of those features have since been added to PowerShell 4.0. Earlier versions of PowerShell do not have many logging features.(Citation: FireEye PowerShell Logging 2016) An organization can gather PowerShell execution details in a data analytic platform to supplement it with other data.', 'filepath': 'https://attack.mitre.org/techniques/T1059/001'})]}"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me Which MITRE technique does this detection belongs to \\\n",
    "powershell EventCode=4104 ScriptBlockText = *Net.webclient* AND ScriptBlockText = *.UploadString* \\\n",
    "| stats count min(_time) as firstTime max(_time) as lastTime by EventCode ScriptBlockText Computer UserID \\\n",
    "| security_content_ctime(firstTime) \\\n",
    "| security_content_ctime(lastTime) \\\n",
    "| windows_exfiltration_over_c2_via_powershell_uploadstring_filter\")\n",
    "res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "73182ba3-3e8d-40fa-a839-94159281fd71",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| filter bysource(All_Traffic.src) | filter bysrcport(All_Traffic.dest_port) | filter by(app=app) | transform value(signature) AS app_name | rename app_name as app_name | rename desc as category | rename desc as signature | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me Which MITRE technique does this detection belongs to | tstats security_content_summariesonly count min(_time) as firstTime max(_time) as lastTime values(All_Traffic.dest_port) as dest_port latest(user) as user from datamodel=Network_Traffic by All_Traffic.src All_Traffic.dest, All_Traffic.app | drop_dm_object_name(All_Traffic) | security_content_ctime(firstTime) | security_content_ctime(lastTime) | lookup remote_access_software remote_appid AS app OUTPUT isutility, description as signature, comment_reference as desc, category | search isutility = True | detect_remote_access_software_usage_traffic_filter ',\n",
       " 'result': '| filter bysource(All_Traffic.src) | filter bysrcport(All_Traffic.dest_port) | filter by(app=app) | transform value(signature) AS app_name | rename app_name as app_name | rename desc as category | rename desc as signature | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc | rename desc as desc',\n",
       " 'source_documents': [Document(page_content='Adversaries may check for Internet connectivity on compromised systems. This may be performed during automated discovery and can be accomplished in numerous ways such as using [Ping](S0097), `tracert`, and GET requests to websites.\\n\\nAdversaries may use the results and responses from these requests to determine if the system is capable of communicating with their C2 servers before attempting to connect to them. The results may also be used to identify routes, redirectors, and proxy servers.', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'System and network discovery techniques normally occur throughout an operation as an adversary learns the environment. Data and events should not be viewed in isolation, but as part of a chain of behavior that could lead to other activities, such as Command and Control, based on the information obtained.\\n\\nMonitor processes and command-line arguments for actions that could be taken to check Internet connectivity.', 'filepath': 'https://attack.mitre.org/techniques/T1016/001'})]}"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me Which MITRE technique does this detection belongs to \\\n",
    "| tstats security_content_summariesonly count min(_time) as firstTime max(_time) as lastTime values(All_Traffic.dest_port) as dest_port latest(user) as user from datamodel=Network_Traffic by All_Traffic.src All_Traffic.dest, All_Traffic.app \\\n",
    "| drop_dm_object_name(All_Traffic) \\\n",
    "| security_content_ctime(firstTime) \\\n",
    "| security_content_ctime(lastTime) \\\n",
    "| lookup remote_access_software remote_appid AS app OUTPUT isutility, description as signature, comment_reference as desc, category \\\n",
    "| search isutility = True \\\n",
    "| detect_remote_access_software_usage_traffic_filter \")\n",
    "res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "7867797f-2ef6-4e2f-b443-64d2f9e5dace",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| filter  \n",
      "| smb_traffic_spike_filter  \n",
      "| smb  \n",
      "| port 139  \n",
      "| port 445  \n",
      "| time span 1h  "
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me Which MITRE technique does this detection belongs to | tstats security_content_summariesonly count from datamodel=Network_Traffic where All_Traffic.dest_port=139 OR All_Traffic.dest_port=445 OR All_Traffic.app=smb by _time span=1h, All_Traffic.src | drop_dm_object_name(All_Traffic) | eventstats max(_time) as maxtime | stats count as num_data_samples max(eval(if(_time >= relative_time(maxtime, -70m@m), count, null))) as count avg(eval(if(_time<relative_time(maxtime, -70m@m), count, null))) as avg stdev(eval(if(_time<relative_time(maxtime, -70m@m), count, null))) as stdev by src | eval upperBound=(avg+stdev*2), isOutlier=if(count > upperBound AND num_data_samples >=50, 1, 0) | where isOutlier=1 | table src count | smb_traffic_spike_filter  ',\n",
       " 'result': '| filter  \\n| smb_traffic_spike_filter  \\n| smb  \\n| port 139  \\n| port 445  \\n| time span 1h  ',\n",
       " 'source_documents': [Document(page_content='PortScan', metadata={'Date': '7/7/17 1:05', 'Source': 'Netflow', 'Subject': '172.16.0.1-192.168.10.50-40620-80-6', 'filepath': 'https://attack.mitre.org/techniques/T1046/'})]}"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me Which MITRE technique does this detection belongs to \\\n",
    "| tstats security_content_summariesonly count from datamodel=Network_Traffic where All_Traffic.dest_port=139 OR All_Traffic.dest_port=445 OR All_Traffic.app=smb by _time span=1h, All_Traffic.src \\\n",
    "| drop_dm_object_name(All_Traffic) \\\n",
    "| eventstats max(_time) as maxtime \\\n",
    "| stats count as num_data_samples max(eval(if(_time >= relative_time(maxtime, -70m@m), count, null))) as count avg(eval(if(_time<relative_time(maxtime, -70m@m), count, null))) as avg stdev(eval(if(_time<relative_time(maxtime, -70m@m), count, null))) as stdev by src \\\n",
    "| eval upperBound=(avg+stdev*2), isOutlier=if(count > upperBound AND num_data_samples >=50, 1, 0) \\\n",
    "| where isOutlier=1 \\\n",
    "| table src count \\\n",
    "| smb_traffic_spike_filter  \")\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "ed24488e-10fc-4ed7-bf2c-cf32d3ce352b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The detection belongs to MITRE technique TP-SA-005 - Powershell Malware. The provided PowerShell code uses the `System.Net.WebClient` class to make an HTTP request, specifically using the `UploadString` method, which suggests that it might be used for exfiltration activities. The output of this detection shows various statistical information related to the usage of this technique on different computers and users, including the first and last time the technique was used on each computer, as well as the exfiltration activity duration. The `security_content_ctime(firstTime)` and `security_content_ctime(lastTime)` functions are used to display the creation time of the relevant security content files, while `windows_exfiltration_over_c2_via_powershell_uploadstring_filter` suggests that the detected activity might be related to exfiltration over a C2 channel using PowerShell's `UploadString` method."
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me Which MITRE technique does this detection belongs to powershell EventCode=4104 ScriptBlockText = *Net.webclient* AND ScriptBlockText = *.UploadString* | stats count min(_time) as firstTime max(_time) as lastTime by EventCode ScriptBlockText Computer UserID | security_content_ctime(firstTime) | security_content_ctime(lastTime) | windows_exfiltration_over_c2_via_powershell_uploadstring_filter',\n",
       " 'result': \"\\nThe detection belongs to MITRE technique TP-SA-005 - Powershell Malware. The provided PowerShell code uses the `System.Net.WebClient` class to make an HTTP request, specifically using the `UploadString` method, which suggests that it might be used for exfiltration activities. The output of this detection shows various statistical information related to the usage of this technique on different computers and users, including the first and last time the technique was used on each computer, as well as the exfiltration activity duration. The `security_content_ctime(firstTime)` and `security_content_ctime(lastTime)` functions are used to display the creation time of the relevant security content files, while `windows_exfiltration_over_c2_via_powershell_uploadstring_filter` suggests that the detected activity might be related to exfiltration over a C2 channel using PowerShell's `UploadString` method.\",\n",
       " 'source_documents': [Document(page_content=\"Adversaries may abuse PowerShell commands and scripts for execution. PowerShell is a powerful interactive command-line interface and scripting environment included in the Windows operating system. (Citation: TechNet PowerShell) Adversaries can use PowerShell to perform a number of actions, including discovery of information and execution of code. Examples include the `Start-Process` cmdlet which can be used to run an executable and the `Invoke-Command` cmdlet which runs a command locally or on a remote computer (though administrator permissions are required to use PowerShell to connect to remote systems).\\n\\nPowerShell may also be used to download and run executables from the Internet, which can be executed from disk or in memory without touching disk.\\n\\nA number of PowerShell-based offensive testing tools are available, including [Empire](S0363),  [PowerSploit](S0194), [PoshC2](S0378), and PSAttack.(Citation: Github PSAttack)\\n\\nPowerShell commands/scripts can also be executed without directly invoking the `powershell.exe` binary through interfaces to PowerShell's underlying `System.Management.Automation` assembly DLL exposed through the .NET framework and Windows Common Language Interface (CLI). (Citation: Sixdub PowerPick Jan 2016)(Citation: SilentBreak Offensive PS Dec 2015)(Citation: Microsoft PSfromCsharp APR 2014)\", metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'If proper execution policy is set, adversaries will likely be able to define their own execution policy if they obtain administrator or system access, either through the Registry or at the command line. This change in policy on a system may be a way to detect malicious use of PowerShell. If PowerShell is not used in an environment, then simply looking for PowerShell execution may detect malicious activity.\\n\\nMonitor for loading and/or execution of artifacts associated with PowerShell specific assemblies, such as System.Management.Automation.dll (especially to unusual process names/locations).(Citation: Sixdub PowerPick Jan 2016)(Citation: SilentBreak Offensive PS Dec 2015)\\n\\nIt is also beneficial to turn on PowerShell logging to gain increased fidelity in what occurs during execution (which is applied to .NET invocations). (Citation: Malware Archaeology PowerShell Cheat Sheet) PowerShell 5.0 introduced enhanced logging capabilities, and some of those features have since been added to PowerShell 4.0. Earlier versions of PowerShell do not have many logging features.(Citation: FireEye PowerShell Logging 2016) An organization can gather PowerShell execution details in a data analytic platform to supplement it with other data.', 'filepath': 'https://attack.mitre.org/techniques/T1059/001'})]}"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me Which MITRE technique does this detection belongs to \\\n",
    "powershell EventCode=4104 ScriptBlockText = *Net.webclient* AND ScriptBlockText = *.UploadString* \\\n",
    "| stats count min(_time) as firstTime max(_time) as lastTime by EventCode ScriptBlockText Computer UserID \\\n",
    "| security_content_ctime(firstTime) \\\n",
    "| security_content_ctime(lastTime) \\\n",
    "| windows_exfiltration_over_c2_via_powershell_uploadstring_filter\")\n",
    "res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "1b363d1b-d7b3-4adf-bd76-1060f6f1f60a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------------\n",
      "This detection belongs to MITRE TECHNIQUE \\_\\_\\_\\_\\_\\_\\_\\_\\_\\_Domain Controller \\_\\_\\_\\_\\_\\_\\_\\_\\_\\_Windows_AD_Rogue_Domain_Controller_Network_Activity_Filter"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'Explain to me Which MITRE technique does this detection belongs to Domain Controllerwindows_ad_rogue_domain_controller_network_activity_filter',\n",
       " 'result': '------------------\\nThis detection belongs to MITRE TECHNIQUE \\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_Domain Controller \\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_\\\\_Windows_AD_Rogue_Domain_Controller_Network_Activity_Filter',\n",
       " 'source_documents': [Document(page_content='Adversaries may check for Internet connectivity on compromised systems. This may be performed during automated discovery and can be accomplished in numerous ways such as using [Ping](S0097), `tracert`, and GET requests to websites.\\n\\nAdversaries may use the results and responses from these requests to determine if the system is capable of communicating with their C2 servers before attempting to connect to them. The results may also be used to identify routes, redirectors, and proxy servers.', metadata={'Date': '4/5/24', 'Source': 'MITRE', 'Subject': 'System and network discovery techniques normally occur throughout an operation as an adversary learns the environment. Data and events should not be viewed in isolation, but as part of a chain of behavior that could lead to other activities, such as Command and Control, based on the information obtained.\\n\\nMonitor processes and command-line arguments for actions that could be taken to check Internet connectivity.', 'filepath': 'https://attack.mitre.org/techniques/T1016/001'})]}"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = qa(\"Explain to me Which MITRE technique does this detection belongs to \\\n",
    "Domain Controller\\\n",
    "windows_ad_rogue_domain_controller_network_activity_filter\\\n",
    "\")\n",
    "res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9cf27892-48f4-4ed3-81ba-2fce502551bf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85a758f0f55946ae8412abc12f1035c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Spotlight running on http://127.0.0.1:56252/'), HBox(children=(Button(description=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#!pip install langchain langchain-openai chromadb renumics-spotlight\n",
    "from renumics import spotlight\n",
    "import pandas as pd\n",
    "#spotlight.show(mitre_sigma)\n",
    "mitre_sigma = pd.read_csv('/Users/Jupiter/Desktop/Python_for_Data_Science/ChatGPT4All_good/spam_data/MITREembed/data/mitre_embeddings (1).csv')\n",
    "spotlight.show(mitre_sigma, dtype={\"image\": spotlight.Image, \"embedding\": spotlight.Embedding})\n",
    "\n",
    "\n"
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}