{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "opened-screen",
   "metadata": {},
   "source": [
    "Sanitizers & Probes\n",
    "==================="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dirty-wholesale",
   "metadata": {},
   "source": [
    "Probe mechanism\n",
    "---------------\n",
    "\n",
    "TritonDSE introduces a probe mechanism. A probe is a class exposing a set of\n",
    "callbacks to register on a given exploration or execution. That enables writing more complex checkers\n",
    "that can be easily registered on a callback manager.\n",
    "\n",
    "A probe as the following interface:\n",
    "\n",
    "```python\n",
    "class ProbeInterface(object):\n",
    "    \"\"\" The Probe interface \"\"\"\n",
    "    def __init__(self):\n",
    "        self._cbs: List[Tuple[CbType, Callable, Optional[str]]] = []\n",
    "\n",
    "    @property\n",
    "    def callbacks(self) -> List[Tuple[CbType, Callable, Optional[str]]]:\n",
    "        return self._cbs\n",
    "\n",
    "    def _add_callback(self, typ: CbType, callback: Callable, arg: str = None):\n",
    "        \"\"\" Add a callback \"\"\"\n",
    "        self._cbs.append((typ, callback, arg))\n",
    "```\n",
    "\n",
    "So to write its own probe one, just have to call `_add_callback` with its own hooks *(usually some other methods)*.\n",
    "They will automatically be picked by the callback manager and registered."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "frequent-fifth",
   "metadata": {},
   "source": [
    "Using built-in sanitizers\n",
    "-------------------------\n",
    "\n",
    "TritonDSE provides few simple sanitizers developped as probes. These sanitizers are the following:\n",
    "\n",
    "* UAFSanitizer: checks for UaF using the simple built-in allocator *(in ProcessState)*\n",
    "* NullDerefSanitizer: checks that no read or write in memory is performed at address 0\n",
    "* FormatStringSanitizer: hooks various libc functions and checks that format string are not controlled\n",
    "\n",
    "The registration of a probe can be done as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "every-bolivia",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tritondse import SymbolicExecutor, Config, Seed, Program, ProcessState, SymbolicExplorator, CoverageStrategy\n",
    "from tritondse.sanitizers import NullDerefSanitizer\n",
    "\n",
    "dse = SymbolicExplorator(Config(), Program(\"crackme_xor\"))\n",
    "\n",
    "dse.callback_manager.register_probe(NullDerefSanitizer())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fitting-assignment",
   "metadata": {},
   "source": [
    "The probe will now be enabled for all executions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ultimate-television",
   "metadata": {},
   "source": [
    "Writing a sanitizer\n",
    "-------------------\n",
    "\n",
    "For the purpose of this tutorial, let's write a sanitizer that will checks `fopen` libc call and its variant `freopen` cannot be hijacked\n",
    "to open an unwanted file *(here /etc/passwd)*. We are first going to check that the string given in input is controllable and if it is,\n",
    "checking by SMT that it can be the intended string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "thirty-density",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tritondse import ProbeInterface, SymbolicExecutor, ProcessState, CbType, SeedStatus\n",
    "from tritondse.types import Addr, SolverStatus\n",
    "\n",
    "class OpenSanitizer(ProbeInterface):\n",
    "    \n",
    "    PASSWD_FILE = \"/etc/passwd\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(OpenSanitizer, self).__init__()\n",
    "        self._add_callback(CbType.PRE_RTN, self.fopen_check, 'fopen')\n",
    "        self._add_callback(CbType.PRE_RTN, self.fopen_check, 'freopen')\n",
    "    \n",
    "    def fopen_check(self, se: SymbolicExecutor, pstate: ProcessState, rtn_name: str, addr: Addr):\n",
    "        # the filepath is located in arg0 (ptr to a string)\n",
    "        string_ptr = se.pstate.get_argument_value(0)\n",
    "        \n",
    "        symbolic = True\n",
    "        cur_ptr = string_ptr\n",
    "        while pstate.read_memory_int(cur_ptr, 1):        # while different from 0\n",
    "            if not se.pstate.is_memory_symbolic(cur_ptr, 1): # check that the byte is symbolic\n",
    "                symbolic = False\n",
    "            cur_ptr += 1\n",
    "        \n",
    "        # if all memory bytes are symbolic and we have enough place to fit our string\n",
    "        if symbolic and (cur_ptr - string_ptr) > len(self.PASSWD_FILE):\n",
    "            \n",
    "            # Try to solve by SMT that the filepath string is /etc/passwd\n",
    "            constraints = [pstate.read_symbolic_memory_byte(string_ptr+i).getAst() == ord(pwd_byte) for i, pwd_byte in enumerate(self.PASSWD_FILE)]\n",
    "            st, model = pstate.solve(constraints)\n",
    "            \n",
    "            if st == SolverStatus.SAT:                          # if formula satisfiable\n",
    "                new_seed = se.mk_new_seed_from_model(model)  # create a new input from the model\n",
    "                new_seed.status = SeedStatus.CRASH           # mark it as crash\n",
    "                se.enqueue_seed(new_seed)                    # enqueue it *(so that it will put in the workspace..)*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pleasant-length",
   "metadata": {},
   "source": [
    "Then we just have to register our new sanitizer!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "organized-newark",
   "metadata": {},
   "outputs": [],
   "source": [
    "dse.callback_manager.register_probe(OpenSanitizer())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "natural-extension",
   "metadata": {},
   "source": [
    "> There are no open in the crackme *(but one can try with another example)* "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neopastis",
   "language": "python",
   "name": "neopastis"
  },
  "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}