1 files changed, 26 insertions, 12 deletions
diff --git a/example/ida/symbol_exec.py b/example/ida/symbol_exec.py
index f019f77d..ffaa9b27 100644
--- a/example/ida/symbol_exec.py
+++ b/example/ida/symbol_exec.py
@@ -34,8 +34,16 @@ class ActionHandlerTranslate(ActionHandler):
 class symbolicexec_t(idaapi.simplecustviewer_t):
 
     def add(self, key, value):
-        self.AddLine("%s = %s" % (expr2colorstr(self.machine.mn.regs.all_regs_ids, key),
-                                  expr2colorstr(self.machine.mn.regs.all_regs_ids, value)))
+        self.AddLine("%s = %s" % (
+            expr2colorstr(
+                key,
+                loc_db=self.loc_db
+            ),
+            expr2colorstr(
+                value,
+                loc_db=self.loc_db
+            )
+        ))
 
     def expand(self, linenum):
         element = self.line2eq[linenum]
@@ -61,11 +69,12 @@ class symbolicexec_t(idaapi.simplecustviewer_t):
         form.Compile()
         form.Execute()
 
-    def Create(self, equations, machine, *args, **kwargs):
+    def Create(self, equations, machine, loc_db, *args, **kwargs):
         if not super(symbolicexec_t, self).Create(*args, **kwargs):
             return False
 
         self.machine = machine
+        self.loc_db = loc_db
         self.line2eq = sorted(equations.items(), key=operator.itemgetter(0))
         self.lines_expanded = set()
 
@@ -119,21 +128,25 @@ def symbolic_exec():
 
     from utils import guess_machine
 
+    start, end = idc.SelStart(), idc.SelEnd()
+
     bs = bin_stream_ida()
-    machine = guess_machine()
+    machine = guess_machine(addr=start)
 
     mdis = machine.dis_engine(bs)
-    start, end = idc.SelStart(), idc.SelEnd()
+
+    if start == idc.BADADDR and end == idc.BADADDR:
+        start = idc.ScreenEA()
+        end = idc.next_head(start) # Get next instruction address
 
     mdis.dont_dis = [end]
-    blocks = mdis.dis_multiblock(start)
-    ira = machine.ira()
-    for block in blocks:
-        ira.add_block(block)
+    asmcfg = mdis.dis_multiblock(start)
+    ira = machine.ira(loc_db=mdis.loc_db)
+    ircfg = ira.new_ircfg_from_asmcfg(asmcfg)
 
     print "Run symbolic execution..."
     sb = SymbolicExecutionEngine(ira, machine.mn.regs.regs_init)
-    sb.run_at(start)
+    sb.run_at(ircfg, start)
     modified = {}
 
     for dst, src in sb.modified(init_state=machine.mn.regs.regs_init):
@@ -141,8 +154,9 @@ def symbolic_exec():
 
     view = symbolicexec_t()
     all_views.append(view)
-    if not view.Create(modified, machine,
-                       "Symbolic Execution - 0x%x to 0x%x" % (start, end)):
+    if not view.Create(modified, machine, mdis.loc_db,
+                       "Symbolic Execution - 0x%x to 0x%x"
+                       % (start, idc.prev_head(end))):
         return
 
     view.Show()