ThreatHuntingNotebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1ceb5156",
   "metadata": {},
   "source": [
    "# Threat Hunting In Rapid Iterations\n",
    "\n",
    "    Author: Alyssa Rahman (@ramen0x3f)\n",
    "    Last Updated: 2022-02-02\n",
    "    \n",
    "The THIRI notebook is designed as a research aide to let you rapidly prototype threat hunting rules."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d64c9c6c",
   "metadata": {},
   "source": [
    "## Getting Started\n",
    "Welcome to the **Threat Hunting In Rapid Iterations (THIRI)** Jupyter notebook! \n",
    "\n",
    "The THIRI notebook is designed as a research aide to let you rapidly prototype threat hunting rules.\n",
    "\n",
    "To get started, click **run**.\n",
    "\n",
    "*Jupyter Pro Tip:*\n",
    "Enable NBExtensions to use the **\"Hide Code Input\"** (eye icon) and **\"Run initialization cells\"** (calculator icon) options above. (Optional, but it will make your output cleaner.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd1c1dd1",
   "metadata": {},
   "source": [
    "### Select Engine\n",
    "Select an engine from the dropdown below. \"Engines\" can be thought of as indicator types or templates. \n",
    "\n",
    "THIRI dynamically loads all sub-classes of Engine() in the utils/engines.py module. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ca5d6e29",
   "metadata": {
    "init_cell": true,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcf9358945564660a15ed1e5f25d54c3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Engine', options=(('CodeCert', <class 'utils.engines.CodeCert'>), ('DLLHijack', <class '…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from inspect import getmembers, isclass\n",
    "from IPython.display import clear_output, Javascript, Markdown\n",
    "from ipywidgets import widgets, VBox, HBox, Label\n",
    "from json import loads\n",
    "from sys import stderr\n",
    "from typing import Union\n",
    "import utils.engines\n",
    "\n",
    "#GLOBAL variables used as gates to keep later cells from executing prematurely\n",
    "GLOBAL = dict.fromkeys([\"selected_engine\", \"engine\", #Engine\n",
    "                        \"inputs\", \"clean_options\", #Engine input/options\n",
    "                        \"edit_boxes\", \"clean_edits\"], None) #Rule edits\n",
    "\n",
    "#Helper function to reset global variables\n",
    "def clear_vars(exceptions = []):\n",
    "    global GLOBAL\n",
    "    #Clearing variables\n",
    "    for x,y in GLOBAL.items():\n",
    "        x = y if x in exceptions else None\n",
    "\n",
    "#Helper function to move to the next cell    \n",
    "def next_cell(e):\n",
    "    display(Javascript('IPython.notebook.select_next()'))\n",
    "    display(Javascript('IPython.notebook.execute_selected_cells()'))\n",
    "    \n",
    "#Helper function called when a new engine is selected\n",
    "def fresh_engine(e):\n",
    "    #Reset vars\n",
    "    clear_vars([\"selected_engine\"])\n",
    "    \n",
    "    #Clear rules from last cell\n",
    "    display(Javascript('IPython.notebook.select(9)'))\n",
    "    display(Javascript('IPython.notebook.execute_selected_cells()'))\n",
    "\n",
    "    #Clear inputs\n",
    "    display(Javascript('IPython.notebook.select(6)'))\n",
    "    display(Javascript('IPython.notebook.execute_selected_cells()'))\n",
    "\n",
    "    #Move to next cell\n",
    "    display(Javascript('IPython.notebook.select(5)'))\n",
    "    display(Javascript('IPython.notebook.execute_selected_cells()'))\n",
    "    \n",
    "#Dynamically build dropdown of classes from utils.engines\n",
    "GLOBAL[\"selected_engine\"] = widgets.Dropdown(\n",
    "    options=[e for e in getmembers(utils.engines, isclass) \n",
    "             if 'utils.engines' in str(e[1]) and e[0] not in ['ABC', 'Engine', 'TemplateClass']],\n",
    "    description='Engine',\n",
    "    disabled=False,\n",
    "    value=None\n",
    ")\n",
    "\n",
    "#Add callback helper and display\n",
    "GLOBAL[\"selected_engine\"].observe(fresh_engine)\n",
    "display(GLOBAL[\"selected_engine\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed677035",
   "metadata": {},
   "source": [
    "### Provide Inputs\n",
    "Each engine has **custom options** and **formats** they support. \n",
    "\n",
    "Some engines (e.g. SerializedObject) may first let you **toggle** select a sub-type of rule to generate. (e.g. inputting a Keyword vs a Chain of Keywords)\n",
    "\n",
    "Fill out the form, and a **prototype rule** will be generated!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "98b096ad",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#Guard condition in case the helper function progresses too fast\n",
    "if GLOBAL[\"selected_engine\"] is not None and GLOBAL[\"selected_engine\"].value is not None:\n",
    "    #Instantiate engine and get first round of inputs\n",
    "    GLOBAL[\"engine\"] = GLOBAL[\"selected_engine\"].value()\n",
    "    opts, tips, descr = GLOBAL[\"engine\"].get_input_type()\n",
    "\n",
    "    #Display engine description\n",
    "    display(Markdown(descr))\n",
    "    \n",
    "    #Move to next cell if no options provided\n",
    "    if opts is None:\n",
    "        search_type = None\n",
    "        next_cell(None)\n",
    "\n",
    "    #Display toggle options if provided\n",
    "    else:    \n",
    "        search_type = widgets.ToggleButtons(\n",
    "            options=opts,\n",
    "            description='Type:',\n",
    "            disabled=False,\n",
    "            tooltips=tips,\n",
    "            value=None\n",
    "        )\n",
    "        \n",
    "        #Add callback helper and display\n",
    "        search_type.observe(next_cell)\n",
    "        display(search_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "65ace61b",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#Guard condition in case the helper function progresses too fast\n",
    "if GLOBAL[\"engine\"] is not None:\n",
    "    #Get custom input/options from engine\n",
    "    #Overwrites search_type.value in case it wasn't provided\n",
    "    st = search_type.value if search_type is not None else \"\"\n",
    "    custom_options = GLOBAL[\"engine\"].get_custom_options(st)\n",
    "\n",
    "    #Display custom options\n",
    "    GLOBAL[\"inputs\"] = []\n",
    "    for opt in custom_options:\n",
    "        desc = Label(opt.description, layout={'width':'10%'})\n",
    "        opt.description = \"\"\n",
    "        \n",
    "        #Can't set Text, Textarea, or SelectMultiple widgets to None\n",
    "        if not isinstance(opt, Union[widgets.Text, widgets.Textarea, widgets.SelectMultiple]):\n",
    "            opt.value = None\n",
    "\n",
    "        #Add calback helper and create box with label\n",
    "        opt.observe(next_cell)\n",
    "        GLOBAL[\"inputs\"].append(HBox([desc, opt]))\n",
    "        \n",
    "    #Put that box inside another box and then mail that box to myself\n",
    "    display(VBox(GLOBAL[\"inputs\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb822019",
   "metadata": {},
   "source": [
    "### Edit Rule\n",
    "All generated rules will be **displayed** below. \n",
    "\n",
    "Make any changes you'd like in the **Edit** box. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e913cbc8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#Guard condition in case the helper function progresses too fast\n",
    "GLOBAL[\"clean_options\"] = {i.children[0].value:i.children[1].value \n",
    "                               for i in GLOBAL[\"inputs\"]} if GLOBAL[\"inputs\"] is not None else {}\n",
    "\n",
    "if GLOBAL[\"engine\"] and len(GLOBAL[\"clean_options\"]['Format']) > 0:\n",
    "    #Generate rule\n",
    "    try: \n",
    "        rules = GLOBAL[\"engine\"].generate(GLOBAL[\"clean_options\"])\n",
    "        \n",
    "        #Iterate\n",
    "        GLOBAL[\"edit_boxes\"] = []\n",
    "        for rulename, rulecontent in rules:\n",
    "            #Create input to edit rule\n",
    "            desc = rulename if rulename is not None else 'unnamed rule'\n",
    "            rule = widgets.Textarea(\n",
    "                disabled=False,\n",
    "                layout={'width': '60%'},\n",
    "                value=rulecontent\n",
    "            )\n",
    "\n",
    "            rule.observe(next_cell)\n",
    "            GLOBAL[\"edit_boxes\"].append(HBox([Label(desc, layout={'width':'30%'}), rule]))\n",
    "\n",
    "        #Put that box inside another box and then mail that box to myself\n",
    "        display(VBox(GLOBAL[\"edit_boxes\"]))\n",
    "    \n",
    "    except RuntimeError as r:\n",
    "        print(f\"[!] {r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd00c80e",
   "metadata": {},
   "source": [
    "### Final Prototype\n",
    "Your final rule prototype will be **printed** below (with any edits automatically rendered). \n",
    "\n",
    "Future plans include automated **publishing** to resources such as VirusTotal, Github, etc. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6ad7037d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Guard condition in case the helper function progresses too fast\n",
    "GLOBAL[\"clean_edits\"] = {e.children[0].value:e.children[1].value \n",
    "                   for e in GLOBAL[\"edit_boxes\"]} if GLOBAL[\"edit_boxes\"] is not None else {}\n",
    "\n",
    "if not any(y is None or len(y) == 0 for x,y in GLOBAL[\"clean_edits\"].items()):\n",
    "    #Just print for now\n",
    "    for rulename, rulecontent in GLOBAL[\"clean_edits\"].items():\n",
    "        print(f\"\\n---------------------------\\n{rulename}\\n---------------------------\\n{rulecontent}\\n\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Initialization Cell",
  "hide_input": true,
  "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.10.1"
  },
  "tags": [
   "hide-input"
  ]
 },
 "nbformat": 4,
 "nbformat_minor": 5
}