luwak-ttp-extractor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "81efc3af",
   "metadata": {},
   "source": [
    "# Luwak TTP Extractor - Extracting TTPs from unstructured threat reports"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ab09cd",
   "metadata": {},
   "source": [
    "## Overview\n",
    "This notebook use [ERNIE 2.0]() English pretrained model to inference MITRE ATTA&CK Enterprise TTPs from unstructured descriptive threat intel data(e.g. blog). It can currently infer common Enterprise Techniques."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26d43847",
   "metadata": {},
   "source": [
    "## Enviroment\n",
    "\n",
    "### Key Requirements\n",
    "\n",
    "Python (64 bit) >= 3.7  \n",
    "pip (64 bit) >= 20.2.2  \n",
    "PaddlePaddle (64 bit) >= 2.0   \n",
    "paddlenlp  \n",
    "nltk\n",
    "\n",
    "Processor Arch x86_64 (arm64 is not supported)\n",
    "\n",
    "### 0. Setup Virtual Env\n",
    "Clone the repository and setup a virtual env with virthenv in Linux:\n",
    "```\n",
    "pip3 install virtualenv\n",
    "mkdir <VENV_NAME>\n",
    "cd <VENV_NAME>\n",
    "virtualenv -p python3 .\n",
    "source ./bin/activate\n",
    "cd <path to this notebook>\n",
    "```\n",
    "\n",
    "### 1. Install Requirements\n",
    "\n",
    "#### CPU\n",
    "\n",
    "**NOT recommmended, time consuming, espacially with large report**\n",
    "\n",
    "Refer to [offical site](https://www.paddlepaddle.org.cn/en/install/quick?docurl=/documentation/docs/en/install/pip/linux-pip_en.html) to install PaddlePaddle.\n",
    "\n",
    "For example, install nltk, paddlenlp and the current stable version 2.4.2 of PaddlePaddle for CPU with pip:\n",
    "```\n",
    "python -m pip install paddlepaddle==2.4.2 -i https://mirror.baidu.com/pypi/simple\n",
    "pip install paddlenlp\n",
    "pip install nltk\n",
    "```\n",
    "\n",
    "#### GPU\n",
    "Refer to [offical site](https://www.paddlepaddle.org.cn/en/install/quick?docurl=/documentation/docs/en/install/pip/linux-pip_en.html) to install PaddlePaddle.\n",
    "\n",
    "For example, install nltk, paddlenlp and the current stable version 2.4.2 of PaddlePaddle for CUDA 11.2 in Linux with pip:\n",
    "```\n",
    "python -m pip install paddlepaddle-gpu==2.4.2.post112 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html\n",
    "pip install paddlenlp\n",
    "pip install nltk\n",
    "```\n",
    "\n",
    "Or you can install requirements from *requirements.txt*. Note that this *requirements.txt* is for PaddlePaddle 2.4.2 with CUDA 11.2 in Linux.\n",
    "```\n",
    "pip install -r requirements.txt\n",
    "```\n",
    "\n",
    "### 2. Download and Merge Model\n",
    "Download nltk punkt and merge pretrained model:\n",
    "```\n",
    "python -c \"import os, nltk; nltk_data_path = os.path.join(os.getcwd(),'nltk_data'); nltk.download('punkt', nltk_data_path);\"\n",
    "\n",
    "./merge_model.sh\n",
    "```\n",
    "\n",
    "Restart the kernel and run the notebook from the **RUN** section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c9258d",
   "metadata": {},
   "source": [
    "## RUN\n",
    "\n",
    "### 1. Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5bd0d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inference\n",
    "from inference import predict_text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03103674",
   "metadata": {},
   "source": [
    "### 2. Infer TTPs from text\n",
    "\n",
    "Put the text to infer TTPs in following variable `text`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d04a883",
   "metadata": {},
   "outputs": [],
   "source": [
    "# text to be predicted\n",
    "text = \"\"\"ACNSHELL is sideloaded by a legitimate executable. \n",
    "It will then create a reverse shell via ncat.exe to the server closed.theworkpc.com\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab3ff827",
   "metadata": {},
   "source": [
    "Call the `predict_text` function to infer related TTPs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0eda0df",
   "metadata": {},
   "outputs": [],
   "source": [
    "o_text, ttps = predict_text(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2567f51",
   "metadata": {},
   "source": [
    "### 3. Result\n",
    "\n",
    "The result include two parts:\n",
    "\n",
    "- `o_text`: The original text.\n",
    "- `ttps`: TTPs inferred from text, list of `dict` objects. Every `dict` contains one sentence from `o_text`, inferred Tactics and Techniques. \n",
    "\n",
    "Example `ttps` for above `o_text`:\n",
    "```\n",
    "[{'sent': 'ACNSHELL is sideloaded by a legitimate executable.',\n",
    "  'tts': [{'tactic_name': 'Persistence',\n",
    "    'tactic_id': 'TA0003',\n",
    "    'technique_name': 'DLLSide-Loading',\n",
    "    'technique_id': 'T1574.002',\n",
    "    'score': 0.9799978}]},\n",
    " {'sent': 'It will then create a reverse shell via ncat.exe to the server closed.theworkpc.com',\n",
    "  'tts': [{'tactic_name': 'Execution',\n",
    "    'tactic_id': 'TA0002',\n",
    "    'technique_name': 'WindowsCommandShell',\n",
    "    'technique_id': 'T1059.003',\n",
    "    'score': 0.98543674}]}]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd07ff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Original text:\")\n",
    "print(o_text)\n",
    "print(\"\\n\")\n",
    "print(\"Predict result:\")\n",
    "for idx, ttp in enumerate(ttps):\n",
    "    sent = ttp['sent']\n",
    "    tts = ttp['tts']\n",
    "    print('idx: {}, sent: {}'.format(idx, sent))\n",
    "    print('tts: {}'.format(tts))\n",
    "    print(\"\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}