From 724db789072da4ae71f58c321771aeef582e99b0 Mon Sep 17 00:00:00 2001
From: Yan <yan@users.noreply.github.com>
Date: Wed, 15 Feb 2017 14:48:55 -0500
Subject: [PATCH] Assign context object to Manticore itself (#13)

* Make Manticore behave like a dict; slim callback
* Update all sample scripts to not take a context
* Add a sample script (count_instructions.py)
* We no longer need to close over Manticore state; add callback directly
* Move Manticore's context to a public attribute
---
 examples/script/count_instructions.py | 29 +++++++++++++++++++++++++++
 examples/script/guide_exec.py         |  2 +-
 examples/script/lads-baby-re.py       |  2 +-
 examples/script/run_callback.py       |  2 +-
 examples/script/state_explore.py      |  2 +-
 examples/script/sym.py                |  2 +-
 manticore/manticore.py                |  9 ++++-----
 7 files changed, 38 insertions(+), 10 deletions(-)
 create mode 100644 examples/script/count_instructions.py

diff --git a/examples/script/count_instructions.py b/examples/script/count_instructions.py
new file mode 100644
index 0000000..1754c25
--- /dev/null
+++ b/examples/script/count_instructions.py
@@ -0,0 +1,29 @@
+
+import sys
+from manticore import Manticore
+
+'''
+Count the number of emulated instructions.
+
+This example uses the context property of the Manticore object to store data
+that's updated by the hook function. Manticore.context is needed to properly
+share data when running with multiple worker processes.
+'''
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        sys.stderr.write("Usage: %s [binary]\n"%(sys.argv[0],))
+        sys.exit(2)
+
+    m = Manticore(sys.argv[1])
+    m.workers = 3
+    m.context['count'] = 0
+
+    def explore(state):
+        m.context['count'] += 1
+
+    m.add_hook(None, explore)
+
+    m.start()
+
+    print "Executed ", m.context['count'], " instructions."
diff --git a/examples/script/guide_exec.py b/examples/script/guide_exec.py
index 9ec02c7..1ff7587 100755
--- a/examples/script/guide_exec.py
+++ b/examples/script/guide_exec.py
@@ -22,7 +22,7 @@ if __name__ == '__main__':
     # Ensure that we ignore all possible branches to libc
     # This hook returns False if we should abandon exploration
     # or True to continue
-    def fork_hook(ctx, state):
+    def fork_hook(state):
         _from, _to = lib.start, lib.start + lib.size
         return not (_from <= state.cpu.PC < _to)
     m.add_fork_hook(fork_hook)
diff --git a/examples/script/lads-baby-re.py b/examples/script/lads-baby-re.py
index 758139e..79e2dc5 100644
--- a/examples/script/lads-baby-re.py
+++ b/examples/script/lads-baby-re.py
@@ -11,7 +11,7 @@ if __name__ == '__main__':
     path = sys.argv[1]
     m = Manticore(path)
 
-    def myhook(ctx, state):
+    def myhook(state):
         flag = ''
         cpu = state.cpu
         arraytop = cpu.R11
diff --git a/examples/script/run_callback.py b/examples/script/run_callback.py
index dbe4be4..e8b4145 100755
--- a/examples/script/run_callback.py
+++ b/examples/script/run_callback.py
@@ -21,7 +21,7 @@ if __name__ == '__main__':
     m = Manticore(None, path, args)
 
     # Trigger an event when PC reaches a certain value
-    def reached_goal(ctx, state):
+    def reached_goal(state):
         cpu = state.cpu
 
         assert cpu.PC == 0x10858
diff --git a/examples/script/state_explore.py b/examples/script/state_explore.py
index 90c7cca..5684ab8 100644
--- a/examples/script/state_explore.py
+++ b/examples/script/state_explore.py
@@ -12,7 +12,7 @@ if __name__ == '__main__':
     # Set to the address of the conditonal checking for the first complex branch
     to_abandon = int(sys.argv[2], 0)
 
-    def explore(ctx, state):
+    def explore(state):
         print "Abandoning state at PC: ", hex(state.cpu.PC)
         state.abandon()
 
diff --git a/examples/script/sym.py b/examples/script/sym.py
index 455e056..6fdccd2 100755
--- a/examples/script/sym.py
+++ b/examples/script/sym.py
@@ -29,7 +29,7 @@ if __name__ == '__main__':
     else:
         target = (0x400a83, 'EBX')
 
-    def entered_func(ctx, state):
+    def entered_func(state):
         '''For ARM, Make R4 symbolic at 0x1082c, as r4 is used in a branch right
            after.
         '''
diff --git a/manticore/manticore.py b/manticore/manticore.py
index f615ef7..5c6db33 100644
--- a/manticore/manticore.py
+++ b/manticore/manticore.py
@@ -170,7 +170,8 @@ class Manticore(object):
         self._verbosity = 0
 
         manager = Manager()
-        self._context = manager.dict()
+
+        self.context = manager.dict()
 
         # XXX(yan) This is a bit obtuse; once PE support is updated this should
         # be refactored out
@@ -225,6 +226,7 @@ class Manticore(object):
 
         self._init_logging()
 
+
     # XXX(yan): args is a temporary hack to include while we continue moving
     # non-Linux platforms to new-style arg handling.
     @property
@@ -316,10 +318,7 @@ class Manticore(object):
         Add a callback to be invoked on executing a program counter. Pass 'None'
         for pc to invoke callback on every instruction.
         '''
-        def _inner(state):
-            callback(self._context, state)
-
-        self._hooks.setdefault(pc, set()).add(_inner)
+        self._hooks.setdefault(pc, set()).add(callback)
 
     def _get_symbol_address(self, symbol):
         '''