2 files changed, 42 insertions, 6 deletions
diff --git a/miasm2/jitter/jitload.py b/miasm2/jitter/jitload.py
index c297ba50..0405b46d 100644
--- a/miasm2/jitter/jitload.py
+++ b/miasm2/jitter/jitload.py
@@ -40,8 +40,10 @@ class CallbackHandler(object):
         self.callbacks = {}  # Key -> [callback list]
 
     def add_callback(self, name, callback):
-        "Add a callback to the key 'name'"
-        self.callbacks[name] = self.callbacks.get(name, []) + [callback]
+        """Add a callback to the key @name, iff the @callback isn't already
+        assigned to it"""
+        if callback not in self.callbacks.get(name, []):
+            self.callbacks[name] = self.callbacks.get(name, []) + [callback]
 
     def set_callback(self, name, *args):
         "Set the list of callback for key 'name'"
@@ -351,3 +353,36 @@ class jitter:
         """Set an unicode string in memory"""
         s = "\x00".join(list(s)) + '\x00' * 3
         self.vm.set_mem(addr, s)
+
+    @staticmethod
+    def handle_lib(jitter):
+        """Resolve the name of the function which cause the handler call. Then
+        call the corresponding handler from users callback.
+        """
+        fname = jitter.libs.fad2cname[jitter.pc]
+        if fname in jitter.user_globals:
+            func = jitter.user_globals[fname]
+        else:
+            log.debug('%s' % repr(fname))
+            raise ValueError('unknown api', hex(jitter.pc), repr(fname))
+        func(jitter)
+        jitter.pc = getattr(jitter.cpu, jitter.ir_arch.pc.name)
+        return True
+
+    def handle_function(self, f_addr):
+        """Add a brakpoint which will trigger the function handler"""
+        self.add_breakpoint(f_addr, self.handle_lib)
+
+    def add_lib_handler(self, libs, user_globals=None):
+        """Add a function to handle libs call with breakpoints
+        @libs: libimp instance
+        @user_globals: dictionnary for defined user function
+        """
+        if user_globals is None:
+            user_globals = {}
+
+        self.libs = libs
+        self.user_globals = user_globals
+
+        for f_addr in libs.fad2cname:
+            self.handle_function(f_addr)
diff --git a/miasm2/jitter/loader/pe.py b/miasm2/jitter/loader/pe.py
index 6b19fc16..0b63583d 100644
--- a/miasm2/jitter/loader/pe.py
+++ b/miasm2/jitter/loader/pe.py
@@ -164,8 +164,9 @@ def vm_load_pe(vm, fdata, align_s=True, load_hdr=True, **kargs):
         # Update min and max addresses
         if min_addr is None or section.addr < min_addr:
             min_addr = section.addr
-        if max_addr is None or section.addr + section.size > max_addr:
-            max_addr = section.addr + max(section.size, len(section.data))
+        max_section_len = max(section.size, len(section.data))
+        if max_addr is None or section.addr + max_section_len > max_addr:
+            max_addr = section.addr + max_section_len
 
     min_addr = pe.rva2virt(min_addr)
     max_addr = pe.rva2virt(max_addr)
@@ -179,8 +180,8 @@ def vm_load_pe(vm, fdata, align_s=True, load_hdr=True, **kargs):
 
     # Copy each sections content in memory
     for section in pe.SHList:
-        log.debug('Map 0x%x bytes to 0x%x' % (len(s.data), pe.rva2virt(s.addr)))
-        vm.set_mem(pe.rva2virt(s.addr), str(s.data))
+        log.debug('Map 0x%x bytes to 0x%x' % (len(section.data), pe.rva2virt(section.addr)))
+        vm.set_mem(pe.rva2virt(section.addr), str(section.data))
 
     return pe