text: {
   "cells": [
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "# -*- coding: utf-8 -*-\n",
      "from miasm.analysis.machine import Machine\n",
      "from miasm.ir.symbexec import SymbolicExecutionEngine\n",
      "from miasm.ir.ir import IRCFG\n",
      "from miasm.expression.expression import LocKey\n",
      "from miasm.arch.x86.regs import *\n",
      "from miasm.core import asmblock\n",
      "from miasm.analysis.binary import Container\n",
      "from future.utils import viewitems\n",
      "from miasm.ir.translators.translator import Translator\n",
      "import networkx as nx\n",
      "import random\n",
      "import z3"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "!unzip -n -P infected vipasana.zip\n",
      "!unzip -n -P infected asprox.zip"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "filename = 'vipasana.bin'\n",
      "target_addr = 0x434DF0\n",
      "#filename = 'asprox.bin'\n",
      "#target_addr = 0x100091AC\n",
      "idc = True"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "def syntax_compare(blk0, blk1):\n",
      "    if len(blk0.lines) != len(blk1.lines):\n",
      "        return False\n",
      "\n",
      "    for l0, l1 in zip(blk0.lines, blk1.lines):\n",
      "        if str(l0)[0] == 'J':\n",
      "            instr0 = str(l0).split(' ')[0]\n",
      "            instr1 = str(l1).split(' ')[0]\n",
      "            if instr0 != instr1:\n",
      "                return False\n",
      "        else:\n",
      "            if str(l0) != str(l1):\n",
      "                return False\n",
      "\n",
      "    return True"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "def semantic_compare(blk0, blk1, ir_arch0, ir_arch1, asmcfg, flag_cmp=False):\n",
      "    src_ircfg = IRCFG(None, ir_arch0.loc_db)\n",
      "    try:\n",
      "        ir_arch0.add_asmblock_to_ircfg(blk0, src_ircfg)\n",
      "    except NotImplementedError:\n",
      "        return False\n",
      "\n",
      "    dst_ircfg = IRCFG(None, ir_arch1.loc_db)\n",
      "    try:\n",
      "        ir_arch1.add_asmblock_to_ircfg(blk1, dst_ircfg)\n",
      "    except NotImplementedError:\n",
      "        return False\n",
      "\n",
      "    if len(src_ircfg.blocks) != len(dst_ircfg.blocks):\n",
      "        return False\n",
      "\n",
      "    for src_lbl, dst_lbl in zip(src_ircfg.blocks, dst_ircfg.blocks):\n",
      "\n",
      "        src_irb = src_ircfg.blocks.get(src_lbl, None)\n",
      "        dst_irb = dst_ircfg.blocks.get(dst_lbl, None)\n",
      "\n",
      "        r = execute_symbolic_execution(\n",
      "                            src_irb, dst_irb, \n",
      "                            ir_arch0, ir_arch1, \n",
      "                            src_ircfg, dst_ircfg,\n",
      "                            flag_cmp)\n",
      "        if r is False:\n",
      "            return False\n",
      "\n",
      "    return True"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "def execute_symbolic_execution(src_irb, dst_irb, \n",
      "                                ir_arch0, ir_arch1, \n",
      "                                src_ircfg, dst_ircfg,\n",
      "                                flag_cmp):\n",
      "\n",
      "    # Ready for Symbolic Execution\n",
      "    src_symbols = {}\n",
      "    dst_symbols = {}\n",
      "\n",
      "    # regs\n",
      "    for i, r in enumerate(all_regs_ids):\n",
      "        src_symbols[r] = all_regs_ids_init[i]\n",
      "        dst_symbols[r] = all_regs_ids_init[i]\n",
      "\n",
      "\n",
      "    # Run symbolic execution\n",
      "    src_sb = SymbolicExecutionEngine(ir_arch0, src_symbols)\n",
      "\n",
      "    for assignblk in src_irb:\n",
      "        skip_update = False\n",
      "        for dst, src in viewitems(assignblk):\n",
      "            if str(dst) in ['EIP', 'IRDst']:\n",
      "                skip_update = True\n",
      "\n",
      "        if not skip_update:\n",
      "            src_sb.eval_updt_assignblk(assignblk)\n",
      "\n",
      "    dst_sb = SymbolicExecutionEngine(ir_arch1, dst_symbols)\n",
      "\n",
      "    for assignblk in dst_irb:\n",
      "        skip_update = False\n",
      "        for dst, src in viewitems(assignblk):\n",
      "            if str(dst) in ['EIP', 'IRDst']:\n",
      "                skip_update = True\n",
      "\n",
      "        if not skip_update:\n",
      "            dst_sb.eval_updt_assignblk(assignblk)\n",
      "\n",
      "    # Equivalence Checking\n",
      "\n",
      "    src_sb.del_mem_above_stack(ir_arch0.sp)\n",
      "    dst_sb.del_mem_above_stack(ir_arch1.sp)\n",
      "\n",
      "    all_memory_ids  = [k for k, v in dst_sb.symbols.memory()] + [k for k, v in src_sb.symbols.memory()]\n",
      "\n",
      "    for k in all_regs_ids + all_memory_ids:\n",
      "\n",
      "        if str(k) == 'EIP':\n",
      "            continue\n",
      "\n",
      "        if not flag_cmp and k in [zf, nf, pf, of, cf, af, df, tf]:\n",
      "            continue\n",
      "\n",
      "        v0 = src_sb.symbols[k]\n",
      "        v1 = dst_sb.symbols[k]\n",
      "\n",
      "        if v0 == v1:\n",
      "            continue\n",
      "\n",
      "        solver = z3.Solver()\n",
      "        try:\n",
      "            z3_r_cond = Translator.to_language('z3').from_expr(v0)\n",
      "        except NotImplementedError:\n",
      "            return False\n",
      "\n",
      "        try:\n",
      "            z3_l_cond = Translator.to_language('z3').from_expr(v1)\n",
      "        except NotImplementedError:\n",
      "            return False\n",
      "\n",
      "        # TODO:\n",
      "        # Please add just one line of code below to check the equivalence of z3_l_cond and z3_r_cond.\n",
      "        # Hint: we can add a constraint to 'solver' with 'add' method, e.g., solver.add(A == B).\n",
      "        # From HERE ------------------------------------------\n",
      "\n",
      "        #solver.add()\n",
      "\n",
      "        # TO HERE ------------------------------------------\n",
      "\n",
      "        r = solver.check()\n",
      "        if r == z3.unsat:\n",
      "            continue\n",
      "\n",
      "        else:\n",
      "            #print(solver.model()) # Counterexample\n",
      "            return False\n",
      "\n",
      "    return True"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "with open(filename, 'rb') as fstream:                                      \n",
      "    cont = Container.from_stream(fstream)\n",
      "    \n",
      "machine = Machine('x86_32')\n",
      "mdis = machine.dis_engine(cont.bin_stream)\n",
      "ir_arch0 = machine.ira(mdis.loc_db)\n",
      "ir_arch1 = machine.ira(mdis.loc_db)\n",
      "\n",
      "asmcfg = mdis.dis_multiblock(target_addr)"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "target_blocks = []\n",
      "for cn, block in enumerate(asmcfg.blocks):\n",
      "    target_blocks.append(block)\n",
      "\n",
      "results = {}\n",
      "\n",
      "for src_blk in target_blocks:\n",
      "    src_ldl = src_blk._loc_key\n",
      "\n",
      "    # Skip a basic block containing only single instruction\n",
      "    if len(src_blk.lines) == 1 and src_blk.lines[0].dstflow():\n",
      "        continue\n",
      "\n",
      "    for dst_blk in target_blocks:\n",
      "        dst_ldl = dst_blk._loc_key\n",
      "\n",
      "        # Skip a basic block containing only single instruction\n",
      "        if len(dst_blk.lines) == 1 and dst_blk.lines[0].dstflow():\n",
      "            continue\n",
      "\n",
      "        if src_ldl == dst_ldl:\n",
      "            continue\n",
      "\n",
      "        if (src_ldl, dst_ldl) in results.keys() or \\\n",
      "            (dst_ldl, src_ldl) in results.keys():\n",
      "            continue\n",
      "        \n",
      "        r_syntax = syntax_compare(src_blk, dst_blk)\n",
      "\n",
      "        if r_syntax:\n",
      "            # If the syntax of two blocks is same, then the semantics of them is also same.\n",
      "            r_semantic = True\n",
      "        else:\n",
      "            # Otherwise, need to compare the semantics of them\n",
      "            r_semantic = semantic_compare(src_blk, dst_blk, ir_arch0, ir_arch1, asmcfg)\n",
      "\n",
      "        results[(src_ldl, dst_ldl)] = [(r_syntax, r_semantic)]"
     ]
    },
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
      "if idc:\n",
      "    G = nx.Graph()\n",
      "    G.add_nodes_from(target_blocks)\n",
      "\n",
      "    for k, v in viewitems(results):\n",
      "        if v[0][0] or v[0][1]:\n",
      "            G.add_edge(k[0], k[1])\n",
      "\n",
      "    # Return a list containing randomlly generated colors\n",
      "    def gen_random_color():\n",
      "\n",
      "        ret = []\n",
      "\n",
      "        r = [x for x in range(256)]\n",
      "        g = [x for x in range(256)]\n",
      "        b = [x for x in range(256)]\n",
      "        random.shuffle(r)\n",
      "        random.shuffle(g)\n",
      "        random.shuffle(b)\n",
      "\n",
      "        for a2, a1, a0 in zip(r,g,b):\n",
      "            color = a2 << 16 | a1 << 8 | a0\n",
      "            ret.append(color)\n",
      "\n",
      "        return ret\n",
      "\n",
      "    random_colors = gen_random_color()\n",
      "    body = ''\n",
      "\n",
      "    for n, conn_nodes in enumerate(nx.connected_components(G)):\n",
      "\n",
      "        if len(conn_nodes) == 1:\n",
      "            continue\n",
      "\n",
      "        for node in conn_nodes: # node is asmblk\n",
      "\n",
      "            if isinstance(node, LocKey):\n",
      "                asmblk = asmcfg.loc_key_to_block(node)\n",
      "                if asmblk:\n",
      "                    for l in asmblk.lines:\n",
      "                        body += 'SetColor(0x%08x, CIC_ITEM, 0x%x);\\n'%(l.offset, random_colors[n])\n",
      "            else:\n",
      "                for l in node.lines:\n",
      "                    body += 'SetColor(0x%08x, CIC_ITEM, 0x%x);\\n'%(l.offset, random_colors[n])\n",
      "        \n",
      "    header = '''\n",
      "#include <idc.idc>\n",
      "static main()\n",
      "{\n",
      "'''\n",
      "    footer = '''\n",
      "}\n",
      "'''\n",
      "\n",
      "    f = open('eq-color.idc', 'w')\n",
      "    f.write(header+body+footer)\n",
      "    f.close()"
     ]
    }
   ],
   "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.6.2"
    }
   },
   "nbformat": 4,
   "nbformat_minor": 4
  }
  