OnWeek2023-langchain-MITRE-rules.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "903a02c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML, display\n",
    "\n",
    "def set_css():\n",
    "  display(HTML('''\n",
    "  <style>\n",
    "    pre {\n",
    "        white-space: pre-wrap;\n",
    "    }\n",
    "  </style>\n",
    "  '''))\n",
    "get_ipython().events.register('pre_run_cell', set_css)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "67a34e6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter OpenAI key: ········\n"
     ]
    }
   ],
   "source": [
    "from getpass import getpass\n",
    "openapi_key = getpass('Enter OpenAI key: ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "7d24ca9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain.document_loaders import JSONLoader, TomlLoader\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "import json\n",
    "import os\n",
    "import zipfile\n",
    "\n",
    "from pathlib import Path\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5a92d36a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "drive_path = \"cti-ATT-CK-v13.0/enterprise-attack/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d5184098",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cti-ATT-CK-v13.0/enterprise-attack/x-mitre-tactic/\n",
      "CPU times: user 448 ms, sys: 13.4 ms, total: 461 ms\n",
      "Wall time: 563 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "tactic_data = []\n",
    "\n",
    "folder_path = drive_path + 'x-mitre-tactic/'\n",
    "\n",
    "print(folder_path)\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "  # print(filename)\n",
    "\n",
    "  description_loader = JSONLoader(\n",
    "    file_path=folder_path+filename,\n",
    "    jq_schema='.objects[] | \" \\(.name) is the name of this tactic. The description of the \\(.name) tactic is: \\(.description). The Mitre Attack code for the \\(.name) tactic is \\(.external_references[].external_id)\"'\n",
    "  )\n",
    "  tactic_data += description_loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3ef355c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cti-ATT-CK-v13.0/enterprise-attack/attack-pattern/\n",
      "CPU times: user 21.8 s, sys: 306 ms, total: 22.1 s\n",
      "Wall time: 22.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Append attack pattern to tactic\n",
    "\n",
    "folder_path = drive_path + 'attack-pattern/'\n",
    "\n",
    "print(folder_path)\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "  # print(filename)\n",
    "\n",
    "  description_loader = JSONLoader(\n",
    "    file_path=folder_path+filename,\n",
    "    jq_schema='.objects[] | \" \\(.name) is the name of this attack pattern. The description of the \\(.name) attack pattern is: \\(.description). The Mitre Attack code for the \\(.name) attack pattern is \\(.external_references[].external_id)\"'\n",
    "  )\n",
    "  # append to tactic\n",
    "  tactic_data += description_loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "085cca69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cti-ATT-CK-v13.0/enterprise-attack/malware/\n",
      "CPU times: user 16.2 s, sys: 232 ms, total: 16.4 s\n",
      "Wall time: 16.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "malware_data = []\n",
    "\n",
    "folder_path = drive_path + 'malware/'\n",
    "\n",
    "print(folder_path)\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "  # print(filename)\n",
    "\n",
    "  description_loader3 = JSONLoader(\n",
    "    file_path=folder_path+filename,\n",
    "    jq_schema='.objects[] | \" \\(.name) is the name of this malware. The description of the \\(.name) malware is: \\(.description). The Mitre Attack code for the \\(.name) malware is \\(.external_references[].external_id)\"'\n",
    "  )\n",
    "  malware_data += description_loader3.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c791b43b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cti-ATT-CK-v13.0/enterprise-attack/course-of-action/\n",
      "CPU times: user 7.7 s, sys: 93.9 ms, total: 7.79 s\n",
      "Wall time: 7.83 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "mitigation_data = []\n",
    "\n",
    "folder_path = drive_path + 'course-of-action/'\n",
    "\n",
    "print(folder_path)\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "  # print(filename)\n",
    "\n",
    "  description_loader2 = JSONLoader(\n",
    "    file_path=folder_path+filename,\n",
    "    jq_schema='.objects[] | \" \\(.name) is the name of this course of action, also known as mitigation. The description of the \\(.name) mitigation is: \\(.description). The Mitre Attack code for the \\(.name) mitigation is \\(.external_references[].external_id)\"'\n",
    "  )\n",
    "  mitigation_data += description_loader2.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aaff5492",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3516"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tactic_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "04ab1e81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rules1/integrations/gcp\n",
      "CPU times: user 12 ms, sys: 4.45 ms, total: 16.4 ms\n",
      "Wall time: 18.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "folder_path = 'rules1/integrations/gcp'\n",
    "\n",
    "print(folder_path)\n",
    "loader = TomlLoader(folder_path)\n",
    "rules_data = loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "3dc157c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.llms import OpenAI\n",
    "import os\n",
    "embeddings = OpenAIEmbeddings(openai_api_key=openapi_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "f6fc9fc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "rules_vStore = Chroma.from_documents(rules_data, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "501df10e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "tactic_vStore = Chroma.from_documents(tactic_data, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3552edac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "malware_vStore = Chroma.from_documents(malware_data, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "257878a6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "mitigation_vStore = Chroma.from_documents(mitigation_data, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "45d6af8f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "llm = OpenAI(temperature=0.5, openai_api_key=openapi_key)\n",
    "chain = load_qa_chain(llm, chain_type=\"stuff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "1d7863d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "' Potential Defense Evasion via PRoot is a rule that identifies the execution of the PRoot utility, an open-source tool for user-space implementation of chroot, mount --bind, and binfmt_misc. Adversaries can leverage this tool to expand the scope of their operations to multiple Linux distributions and simplify their necessary efforts. It also provides emulation capabilities that allow for malware built on other architectures, such as ARM, to be run. The post-exploitation technique called bring your own filesystem (BYOF) can be used by the threat actors to execute malicious payload or elevate privileges or perform network scans or orchestrate another attack.'"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"Describe Potential Defense Evasion via PRoot\"\n",
    "docs = rules_vStore.similarity_search(query, include_metadata=True, k=2)\n",
    "chain.run(input_documents=docs, question=query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "0a376ded",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Document(page_content=' DLL Side-Loading is the name of this attack pattern. The description of the DLL Side-Loading attack pattern is: Programs may specify DLLs that are loaded at runtime. Programs that improperly or vaguely specify a required DLL may be open to a vulnerability in which an unintended DLL is loaded. Side-loading vulnerabilities specifically occur when Windows Side-by-Side (WinSxS) manifests (Citation: MSDN Manifests) are not explicit enough about characteristics of the DLL to be loaded. Adversaries may take advantage of a legitimate program that is vulnerable to side-loading to load a malicious DLL. (Citation: Stewart 2014)\\n\\nAdversaries likely use this technique as a means of masking actions they perform under a legitimate, trusted system or software process.. The Mitre Attack code for the DLL Side-Loading attack pattern is null', metadata={'source': '/Users/kirtikirti/langchain/cti-ATT-CK-v13.0/enterprise-attack/attack-pattern/attack-pattern--b2001907-166b-4d71-bb3c-9d26c871de09.json', 'seq_num': 3}),\n",
      " Document(page_content=' DLL Side-Loading is the name of this attack pattern. The description of the DLL Side-Loading attack pattern is: Programs may specify DLLs that are loaded at runtime. Programs that improperly or vaguely specify a required DLL may be open to a vulnerability in which an unintended DLL is loaded. Side-loading vulnerabilities specifically occur when Windows Side-by-Side (WinSxS) manifests (Citation: MSDN Manifests) are not explicit enough about characteristics of the DLL to be loaded. Adversaries may take advantage of a legitimate program that is vulnerable to side-loading to load a malicious DLL. (Citation: Stewart 2014)\\n\\nAdversaries likely use this technique as a means of masking actions they perform under a legitimate, trusted system or software process.. The Mitre Attack code for the DLL Side-Loading attack pattern is null', metadata={'source': '/Users/kirtikirti/langchain/cti-ATT-CK-v13.0/enterprise-attack/attack-pattern/attack-pattern--b2001907-166b-4d71-bb3c-9d26c871de09.json', 'seq_num': 4}),\n",
      " Document(page_content=' DLL Side-Loading is the name of this attack pattern. The description of the DLL Side-Loading attack pattern is: Programs may specify DLLs that are loaded at runtime. Programs that improperly or vaguely specify a required DLL may be open to a vulnerability in which an unintended DLL is loaded. Side-loading vulnerabilities specifically occur when Windows Side-by-Side (WinSxS) manifests (Citation: MSDN Manifests) are not explicit enough about characteristics of the DLL to be loaded. Adversaries may take advantage of a legitimate program that is vulnerable to side-loading to load a malicious DLL. (Citation: Stewart 2014)\\n\\nAdversaries likely use this technique as a means of masking actions they perform under a legitimate, trusted system or software process.. The Mitre Attack code for the DLL Side-Loading attack pattern is T1073', metadata={'source': '/Users/kirtikirti/langchain/cti-ATT-CK-v13.0/enterprise-attack/attack-pattern/attack-pattern--b2001907-166b-4d71-bb3c-9d26c871de09.json', 'seq_num': 1}),\n",
      " Document(page_content=' DLL Side-Loading is the name of this attack pattern. The description of the DLL Side-Loading attack pattern is: Programs may specify DLLs that are loaded at runtime. Programs that improperly or vaguely specify a required DLL may be open to a vulnerability in which an unintended DLL is loaded. Side-loading vulnerabilities specifically occur when Windows Side-by-Side (WinSxS) manifests (Citation: MSDN Manifests) are not explicit enough about characteristics of the DLL to be loaded. Adversaries may take advantage of a legitimate program that is vulnerable to side-loading to load a malicious DLL. (Citation: Stewart 2014)\\n\\nAdversaries likely use this technique as a means of masking actions they perform under a legitimate, trusted system or software process.. The Mitre Attack code for the DLL Side-Loading attack pattern is CAPEC-641', metadata={'source': '/Users/kirtikirti/langchain/cti-ATT-CK-v13.0/enterprise-attack/attack-pattern/attack-pattern--b2001907-166b-4d71-bb3c-9d26c871de09.json', 'seq_num': 2})]\n"
     ]
    }
   ],
   "source": [
    "test_query = \"What is lateral movement?\"\n",
    "pprint(tactic_vStore.similarity_search(query=test_query, t=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88fec8ad",
   "metadata": {},
   "source": [
    "## MVP Prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "c80250a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "tactic_qa = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=tactic_vStore.as_retriever(k=6, verbose=True))\n",
    "malware_qa = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=malware_vStore.as_retriever(k=6))\n",
    "mitigation_qa = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=mitigation_vStore.as_retriever(k=9))\n",
    "rules_qa = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=rules_vStore.as_retriever(k=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "6bc0e8aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "def run_queries(query):\n",
    "    description = malware_qa.run(f\"For the malware, '{scenario}', what is the definition?\")\n",
    "    summary = malware_qa.run(summary_template)\n",
    "    number = malware_qa.run(f\"What is the Mitre Attack code for the {mitre_type}, {scenario}? Return only this code.\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "73b53086",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def detection_rules(rule_name):\n",
    "    #summary_template = rulesf\"Describe {rule_name}. What is it's risk score and severity?\"\n",
    "    description = rules_qa.run(f\"Describe {rule_name}.\")\n",
    "    techniques = rules_qa.run(f\"What attack techniques are used for {rule_name}.\")\n",
    "    severity = rules_qa.run(f\"What is the risk score and severity of {rule_name}.\")\n",
    "    mitigations = rules_qa.run(f\"Give some mitigations for {rule_name}.\")\n",
    "    print (\n",
    "    f\"{description}\\n\\n\", f\"Attack techniques for {rule_name}: {techniques} \\n\\n\", f\"Severity of {rule_name}:{severity} \\n\\n\", f\"Possible mitigations of {rule_name}: {mitigations}\",\n",
    "    \"\\n\\n\\n\\n\\n\",\n",
    "      \"Thanks for testing the WIP Elastic mitre buddy!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "76c059a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def mitre_buddy(mitre_type, scenario):\n",
    "    summary_template = f\"For the {mitre_type}, {scenario}, can you summarize the attack techniques or malware techniques from the {mitre_type}'s definition? Keep it brief, keywords only.\"\n",
    "    if mitre_type == \"malware\":\n",
    "        description = malware_qa.run(f\"For the malware, '{scenario}', what is the definition?\")\n",
    "        summary = malware_qa.run(summary_template)\n",
    "        number = malware_qa.run(f\"What is the Mitre Attack code for the {mitre_type}, {scenario}? Return only this code.\")    \n",
    "    elif mitre_type == 'tactic':\n",
    "        description = tactic_qa.run(f\"For the tactic, '{scenario}', what is the definition?\")\n",
    "        summary = tactic_qa.run(summary_template)\n",
    "        number = tactic_qa.run(f\"What is the Mitre Attack code for the {mitre_type}, {scenario}? Return only this code.\")\n",
    "    elif mitre_type == 'attack-pattern':\n",
    "        description = tactic_qa.run(f\"For the attack pattern, '{scenario}', what is the definition?\")\n",
    "        summary = tactic_qa.run(summary_template)\n",
    "        number = tactic_qa.run(f\"What is the Mitre Attack code for the {mitre_type}, {scenario}? Return only this code.\")     \n",
    "    else:\n",
    "        description = \"Custom answer: No matches.\"\n",
    "    mitigations = mitigation_qa.run(f\"What are the mitigations to protect against {summary}? Please list out the corresponding Mitre Attack codes of the mitigations.\")\n",
    "    if mitigations == \"I don't know.\":\n",
    "        mitigations = \"Sorry, but I don't know the answer! Go to the docs and tell the LangChain team!\"\n",
    "    print(\n",
    "      f\"{description} uses or is related to the keywords, {summary}\",\n",
    "      f\"The code for {scenario}, if found, is {number}.\",\n",
    "      f\"You asked how to mitigate the {mitre_type}, {scenario}. Here's how: \\n\\n\",\n",
    "      mitigations,\n",
    "      \"\\n\\n\\n\\n\\n\",\n",
    "      \"Thanks for testing the WIP Elastic mitre buddy!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "69d5d714",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "23"
      ]
     },
     "execution_count": 150,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(rules_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f058e43",
   "metadata": {},
   "source": [
    "# Elastic Security Buddy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e5e93a8",
   "metadata": {},
   "source": [
    "## Demo Area"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "44833333",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "defense evasion\n"
     ]
    }
   ],
   "source": [
    "rule_name = input()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "880ecf81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Defense evasion is a tactic used by adversaries to avoid detection and maintain access to a system. It can involve techniques such as hiding files and directories, masquerading, path interception, and renaming system utilities.\n",
      "\n",
      " Attack techniques for defense evasion:  T1211 - Exploitation for Defense Evasion and T1574 - Hijack Execution Flow. \n",
      "\n",
      " Severity of defense evasion: I don't know. \n",
      "\n",
      " Possible mitigations of defense evasion:  Mitigations for defense evasion include: monitoring system configurations for unexpected changes, maintaining least privilege access, and using application whitelisting to prevent unauthorized programs from running. Additionally, implementing endpoint detection and response (EDR) solutions can help detect and respond to defense evasion tactics. \n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Thanks for testing the WIP Elastic mitre buddy!\n"
     ]
    }
   ],
   "source": [
    "detection_rules(rule_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "b94c248d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GCP Service Account Creation\n"
     ]
    }
   ],
   "source": [
    "rule_name = input()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "453af9c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " GCP Service Account Creation is a rule that identifies when a new service account is created in Google Cloud Platform (GCP). A service account is a special type of account used by an application or a virtual machine (VM) instance, not a person. Applications use service accounts to make authorized API calls, authorized as either the service account itself, or as G Suite or Cloud Identity users through domain-wide delegation. If service accounts are not tracked and managed properly, they can present a security risk. An adversary may create a new service account to use during their operations in order to avoid using a standard user account and attempt to evade detection.\n",
      "\n",
      " Attack techniques for GCP Service Account Creation:  The attack technique used for GCP Service Account Creation is T1136, Create Account. \n",
      "\n",
      " Severity of GCP Service Account Creation: The risk score of GCP Service Account Creation is 21 and the severity is low. \n",
      "\n",
      " Possible mitigations of GCP Service Account Creation:  Verify that the behavior was expected and add exceptions to filter expected behavior. Additionally, track and manage service accounts properly to reduce the risk of abuse. \n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Thanks for testing the WIP Elastic mitre buddy!\n"
     ]
    }
   ],
   "source": [
    "detection_rules(rule_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e796c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mitre_type = input() # [\"malware\", \"tactic\", \"attack-pattern\"]\n",
    "scenario = input()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "ac0f5619",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <style>\n",
       "    pre {\n",
       "        white-space: pre-wrap;\n",
       "    }\n",
       "  </style>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " The Credential Access tactic consists of techniques for stealing credentials like account names and passwords. Techniques used to get credentials include keylogging or credential dumping. Using legitimate credentials can give adversaries access to systems, make them harder to detect, and provide the opportunity to create more accounts to help achieve their goals. uses or is related to the keywords,  Keylogging, credential dumping. The code for Credential Access, if found, is  T1589.001. You asked how to mitigate the tactic, Credential Access. Here's how: \n",
      "\n",
      "  The mitigations to protect against Keylogging and credential dumping are Credential Dumping Mitigation (T1003), Valid Accounts Mitigation (T1078), and Credentials in Registry Mitigation (T1214). \n",
      "\n",
      "\n",
      "\n",
      "\n",
      " Thanks for testing the WIP Elastic mitre buddy!\n"
     ]
    }
   ],
   "source": [
    "mitre_buddy(mitre_type, scenario)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c099d67",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.10.4",
   "language": "python",
   "name": "project_env"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}