Add hook decorator (#28)

* Add m.add_hook test

* Add @m.hook test

* Add `hook` decorator for convenience

* Update readme and examples

* Update run_callback

* Improve `add_hook` docstring

expound on callback structure

* Rm debug print

* Improve docstring

README.md

@@ -66,6 +66,7 @@ hook_pc = 0x400ca0
 
 m = Manticore('./path/to/binary')
 
+@m.hook(hook_pc)
 def hook(state):
   cpu = state.cpu
   print 'eax', cpu.EAX
@@ -73,6 +74,5 @@ def hook(state):
   
   m.terminate()  # tell Manticore to stop
 
-m.add_hook(hook_pc, hook)
 m.run()
 ```

examples/script/count_instructions.py

@@ -19,11 +19,10 @@ if __name__ == '__main__':
     m.workers = 3
     m.context['count'] = 0
 
+    @m.hook(None)
     def explore(state):
         m.context['count'] += 1
 
-    m.add_hook(None, explore)
-
     m.run()
 
     print "Executed ", m.context['count'], " instructions."

examples/script/lads-baby-re.py

@@ -11,6 +11,7 @@ if __name__ == '__main__':
     path = sys.argv[1]
     m = Manticore(path)
 
+    @m.hook(0x109f0)
     def myhook(state):
         flag = ''
         cpu = state.cpu
@@ -24,7 +25,5 @@ if __name__ == '__main__':
         print 'flag is:', flag
         m.terminate()
 
-    m.add_hook(0x109f0, myhook)
-
     m.run()
     print 'done'

examples/script/run_callback.py

@@ -5,34 +5,24 @@ from manticore import Manticore
 
 # This example demonstrates a basic hook (PC register)
 
-def get_args():
-    class Args(object): pass
-    args = Args()
-    args.replay = None; args.data = ''; args.dumpafter = 0; args.maxstates = 0;
-    args.maxstorage = 0; args.stats = True; args.verbose = False; args.log = '-';
-    return args
-
 if __name__ == '__main__':
     path = sys.argv[1]
-    args = get_args()
+    pc = int(sys.argv[2], 0)
 
-    args.programs = sys.argv[1:]
-    # Create a new Manticore object
-    m = Manticore(None, path, args)
+    m = Manticore(path)
 
     # Trigger an event when PC reaches a certain value
+    @m.hook(pc)
     def reached_goal(state):
         cpu = state.cpu
 
-        assert cpu.PC == 0x10858
+        assert cpu.PC == pc
 
-        instruction = cpu.read(cpu.PC, 4)
+        instruction = cpu.read_int(cpu.PC)
         print "Execution goal reached."
-        print "Instruction bytes: {:08x}".format(cpu.pc)
+        print "Instruction bytes: {:08x}".format(instruction)
 
-    m.add_pc_hook(0x10858, reached_goal)
-
-    # Start path exploration. start() returns when Manticore
+    # Start path exploration. m.run() returns when Manticore
     # finishes
     m.run()
 

examples/script/state_explore.py

@@ -12,11 +12,11 @@ if __name__ == '__main__':
     # Set to the address of the conditonal checking for the first complex branch
     to_abandon = int(sys.argv[2], 0)
 
+    @m.hook(to_abandon)
     def explore(state):
         print "Abandoning state at PC: ", hex(state.cpu.PC)
         state.abandon()
 
     print "Adding hook to: ", hex(to_abandon)
-    m.add_hook(to_abandon, explore)
 
     m.run()

manticore/manticore.py

@@ -305,10 +305,21 @@ class Manticore(object):
             logging.getLogger(log_type).setLevel(level)
             self._verbosity = setting
 
+    def hook(self, pc):
+        '''
+        A decorator used to register a hook function for a given instruction address
+        (`pc`). Equivalent to calling `add_hook`.
+        '''
+        def decorator(f):
+            self.add_hook(pc, f)
+            return f
+        return decorator
+
     def add_hook(self, pc, callback):
         '''
         Add a callback to be invoked on executing a program counter. Pass 'None'
-        for pc to invoke callback on every instruction.
+        for pc to invoke callback on every instruction. `callback` should be a callable
+        that takes one `manticore.core.executor.State` argument.
         '''
         if not (isinstance(pc, (int, long)) or pc is None):
             raise TypeError("pc must be either an int or None, not {}".format(pc.__class__.__name__))

test/test_manticore.py

@@ -0,0 +1,21 @@
+import unittest
+
+from manticore import Manticore
+
+class ManticoreTest(unittest.TestCase):
+    def setUp(self):
+        self.m = Manticore('test/binaries/arguments_linux_amd64')
+
+    def test_add_hook(self):
+        def tmp(state):
+            pass
+        entry = 0x00400e40
+        self.m.add_hook(entry, tmp)
+        self.assertTrue(tmp in self.m._hooks[entry])
+
+    def test_hook_dec(self):
+        entry = 0x00400e40
+        @self.m.hook(entry)
+        def tmp(state):
+            pass
+        self.assertTrue(tmp in self.m._hooks[entry])