Kibana_MMBot_Triage_Example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Triage Demonstration Using Elasticsearch / Kibana "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some helper functions below needed to insert office document prediction results into a local Elasticsearch instance.  This was tested with ES / Kibana 5.1.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import mmbot as mmb\n",
    "from elasticsearch import Elasticsearch\n",
    "import time\n",
    "import requests\n",
    "import json\n",
    "\n",
    "def wipeESCreateMappings(myindex, doc_type):\n",
    "    '''\n",
    "    Helper function to destroy all data in myindex if it exists and recreate field mappings\n",
    "    '''\n",
    "    es = Elasticsearch([{'host':'localhost', 'port':9200}])\n",
    "    try:\n",
    "        es.indices.delete(myindex, ignore=400)\n",
    "    except:\n",
    "        print \"Mapping does not exist\"\n",
    "    es.indices.get_settings()\n",
    "    newmapping = {\n",
    "        doc_type: {\n",
    "            \"properties\" : {\n",
    "                \"prediction\":  {\"type\":\"string\", \"index\":\"not_analyzed\"}, \n",
    "                \"md5\":  {\"type\":\"string\", \"index\":\"not_analyzed\"},\n",
    "                \"featureprint\":  {\"type\":\"string\", \"index\":\"not_analyzed\"},\n",
    "                \"scan_date\": {\"type\": \"date\", \"format\":\"YYYY-MM-dd\"},\n",
    "                \"extracted_vba\": {\"type\":\"string\"},\n",
    "                \"function_names\": {\"type\":\"string\"},\n",
    "                \"sample_tag\": {\"type\":\"string\"},\n",
    "                \"vba_avg_param_per_func\": {\"type\": \"float\"},\n",
    "                \"vba_cnt_comment_loc_ratio\": {\"type\": \"float\"},\n",
    "                \"vba_cnt_comments\": {\"type\": \"long\"},\n",
    "                \"vba_cnt_func_loc_ratio\": {\"type\": \"float\"},\n",
    "                \"vba_cnt_functions\": {\"type\": \"long\"},\n",
    "                \"vba_cnt_loc\": {\"type\": \"long\"},\n",
    "                \"vba_entropy_chars\": {\"type\": \"float\"},\n",
    "                \"vba_entropy_func_names\": {\"type\": \"float\"},\n",
    "                \"vba_entropy_words\": {\"type\": \"float\"},\n",
    "                \"vba_mean_loc_per_func\": {\"type\": \"float\"}\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    es = Elasticsearch([{'host':'localhost', 'port':9200}])\n",
    "    es.indices.create(index=myindex, ignore=400)\n",
    "    put_mapping_result = es.indices.put_mapping(index=myindex, doc_type=doc_type, ignore=400, body=newmapping)\n",
    "    print put_mapping_result \n",
    "    es.indices.get_settings()\n",
    "\n",
    "def insertIntoES(doc, myindex, doc_type):\n",
    "    '''\n",
    "    Inserts the dictionary provided into the specified index and doc_type\n",
    "    '''\n",
    "    res = requests.get('http://localhost:9200/')\n",
    "    es = Elasticsearch([{'host':'localhost', 'port':9200}])\n",
    "    es.index(index=myindex, doc_type=doc_type, id=doc['md5'], body=json.dumps(doc))\n",
    "\n",
    "def insertResultIntoES(prediction_result, sample_tag=None):\n",
    "    '''\n",
    "    Takes a prediction dataframe (from calling mm_predict) and parses out specific fields\n",
    "    then inserts into ES for visualization\n",
    "    '''\n",
    "    for i in range(len(prediction_result)):\n",
    "        vba_clean = prediction_result.iloc[i].extracted_vba[0:5000].encode(\"utf-8\")\n",
    "        scan_date = time.strftime('%Y-%m-%d')\n",
    "        doc = {'md5':prediction_result.iloc[i].md5,\n",
    "            'prediction':prediction_result.iloc[i].prediction,\n",
    "            'featureprint':prediction_result.iloc[i].featureprint,\n",
    "            'vba_cnt_comments':prediction_result.iloc[i].vba_cnt_comments,\n",
    "            'vba_cnt_functions':prediction_result.iloc[i].vba_cnt_functions,\n",
    "            'vba_cnt_loc':prediction_result.iloc[i].vba_cnt_loc,\n",
    "            'vba_entropy_chars':prediction_result.iloc[i].vba_entropy_chars,\n",
    "            'vba_entropy_func_names':prediction_result.iloc[i].vba_entropy_func_names,\n",
    "            'vba_entropy_words':prediction_result.iloc[i].vba_entropy_words,\n",
    "            'vba_mean_loc_per_func':prediction_result.iloc[i].vba_mean_loc_per_func,\n",
    "            'extracted_vba':vba_clean,\n",
    "            'scan_date':scan_date\n",
    "        }        \n",
    "        \n",
    "        if sample_tag is not None:\n",
    "            doc['sample_tag'] = sample_tag\n",
    "        insertIntoES(doc, 'triage', 'triage_sample')\n",
    "    print \"Completed inserting {} documents\".format(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Re)Create field mapping in ES, make predictions, and insert into ES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{u'acknowledged': True}\n",
      "Completed inserting 20 documents\n",
      "CPU times: user 3.71 s, sys: 1.29 s, total: 5 s\n",
      "Wall time: 5.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "wipeESCreateMappings('triage', 'triage_sample')\n",
    "\n",
    "files_to_triage = '/home/test/vt_all/triage_demo/'\n",
    "mymacrobot = mmb.MaliciousMacroBot()\n",
    "mymacrobot.mmb_init_model(modelRebuild=False)\n",
    "result = mymacrobot.mmb_predict(files_to_triage, datatype='filepath')\n",
    "insertResultIntoES(result, sample_tag=\"triage\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}