{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from collections import defaultdict\n",
    "\n",
    "def approximate_taint_tracking(segment_size, input_partitions, output_partitions):\n",
    "    input_segments = defaultdict(set)  # map from input segment to input partition indexes\n",
    "    output_segments = defaultdict(set)  # map from output segment to output partition indexes\n",
    "    partition_indexes_map = defaultdict(set)  # map from input partition index to output partition indexes\n",
    "    tainted_partitions = []  # list of TAINTED output partition and the input partition they are tainted by\n",
    "    \n",
    "    # build segments for input partitions\n",
    "    for partition_index, partition in enumerate(input_partitions):\n",
    "        for i in range(len(partition) - segment_size):\n",
    "            input_segments[partition[i:i + segment_size]].add(partition_index)\n",
    "\n",
    "\n",
    "    # build segments for output partitions\n",
    "    for partition_index, partition in enumerate(output_partitions):\n",
    "        for i in range(len(partition) - segment_size):\n",
    "            output_segments[partition[i:i + segment_size]].add(partition_index)\n",
    "\n",
    "    # map input segments to output segments\n",
    "    for segment in input_segments:\n",
    "        for partition_index in input_segments[segment]:\n",
    "            partition_indexes_map[partition_index].update(output_segments[segment])\n",
    "\n",
    "    # mark input partitions that are TAINTED by multipule output_partitions as injections\n",
    "    for partition_index in partition_indexes_map:\n",
    "        tainted_output_partitions = [output_partitions[output_partition_index] \n",
    "                                     for output_partition_index \n",
    "                                     in partition_indexes_map[partition_index]]\n",
    "        \n",
    "        # add tainted_output_partitions to tainted_partitions\n",
    "        if tainted_output_partitions:\n",
    "            tainted_partitions.append([\n",
    "                # if more the one output is TAINTED by a single input control was INJECTED\n",
    "                'INJECTED' if len(tainted_output_partitions) > 1 else 'TAINTED', \n",
    "                input_partitions[partition_index], \n",
    "                tainted_output_partitions,\n",
    "            ])\n",
    "    \n",
    "    return tainted_partitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Example\n",
    "==\n",
    "\n",
    "JSON\n",
    "--\n",
    "\n",
    "raw input\n",
    "\n",
    "```JSON\n",
    "{\"email\": \"me@example.com\", \"password\": \"mYw0rd!3 OR '1'='1'\"}\n",
    "```\n",
    "\n",
    "partioned input\n",
    "\n",
    "```python\n",
    "['me@example.com', 'mYw0rd!3 OR \\'1\\'=\\'1\\'']\n",
    "```\n",
    "SQL\n",
    "--\n",
    "\n",
    "raw output\n",
    "\n",
    "```SQL\n",
    "SELECT user_id\n",
    "FROM users\n",
    "WHERE users_email = me@example.com\n",
    "AND Users_pw = mYw0rd!3\n",
    "OR '1'='1';\n",
    "```\n",
    "\n",
    "partioned output\n",
    "\n",
    "```python\n",
    "['user_id', 'users', 'users_email = me@example.com', 'Users_pw = mYw0rd!3', '\\'1\\'=\\'1\\'']\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "json_partitions = ['me@example.com', 'mYw0rd!3 OR \\'1\\'=\\'1\\'']\n",
    "\n",
    "sql_partitions = ['user_id', 'users', 'users_email = me@example.com', 'Users_pw = mYw0rd!3', '\\'1\\'=\\'1\\'']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Detecting both TAINTED and INJECTED outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "pprint(approximate_taint_tracking(segment_size=5, input_partitions=json_partitions, output_partitions=sql_partitions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Tuning the false positive false negative rate with segment sizing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for i in range(1,15):\n",
    "    print \"\\n\\nsegment_size={}\\n-\".format(i)\n",
    "    pprint(approximate_taint_tracking(\n",
    "        segment_size=i, \n",
    "        input_partitions=json_partitions, \n",
    "        output_partitions=sql_partitions,\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}