Merge pull request #26 from commial/fix-vm_loads

Fix vm loads
Everything seems ok, thanks
Merry Christmiasm to you !

8 files changed, 129 insertions, 102 deletions

diff --git a/example/test_jit_arm.py b/example/test_jit_arm.py
index e4eddad3..12878a30 100644
--- a/example/test_jit_arm.py
+++ b/example/test_jit_arm.py
@@ -3,7 +3,6 @@
 from miasm2.analysis import debugging, gdbserver
 
 from miasm2.analysis.sandbox import Sandbox_Linux_arml
-from miasm2.jitter.loader.elf import vm_load_elf, preload_elf
 from miasm2.analysis.machine import Machine
 import logging
 
diff --git a/example/test_jit_arm_sc.py b/example/test_jit_arm_sc.py
index 2cac8686..0ec2a5ee 100644
--- a/example/test_jit_arm_sc.py
+++ b/example/test_jit_arm_sc.py
@@ -4,7 +4,6 @@ from miasm2.analysis import debugging, gdbserver
 
 from miasm2.analysis.sandbox import Sandbox_Linux_armb_str
 from miasm2.analysis.sandbox import Sandbox_Linux_arml_str
-from miasm2.jitter.loader.elf import vm_load_elf, preload_elf
 from miasm2.analysis.machine import Machine
 from elfesteem.strpatchwork import StrPatchwork
 import logging
diff --git a/example/test_jit_mips32.py b/example/test_jit_mips32.py
index f09e2d81..d027ec50 100644
--- a/example/test_jit_mips32.py
+++ b/example/test_jit_mips32.py
@@ -3,9 +3,8 @@
 from argparse import ArgumentParser
 from miasm2.analysis import debugging, gdbserver
 from miasm2.jitter.csts import *
-
-from miasm2.jitter.loader.elf import vm_load_elf, preload_elf
 from miasm2.analysis.machine import Machine
+
 from pdb import pm
 
 parser = ArgumentParser(
diff --git a/example/test_jit_msp430.py b/example/test_jit_msp430.py
index 4df9a869..5aa1f340 100644
--- a/example/test_jit_msp430.py
+++ b/example/test_jit_msp430.py
@@ -3,8 +3,6 @@
 from argparse import ArgumentParser
 from miasm2.analysis import debugging, gdbserver
 from miasm2.jitter.csts import *
-
-from miasm2.jitter.loader.elf import vm_load_elf, preload_elf
 from miasm2.analysis.machine import Machine
 
 parser = ArgumentParser(
diff --git a/miasm2/analysis/binary.py b/miasm2/analysis/binary.py
index a46a252c..17c87761 100644
--- a/miasm2/analysis/binary.py
+++ b/miasm2/analysis/binary.py
@@ -122,7 +122,7 @@ class ContainerPE(Container):
         # Build executable instance
         try:
             if vm is not None:
-                self._executable = vm_load_pe(vm, filename)
+                self._executable = vm_load_pe(vm, data)
             else:
                 self._executable = pe_init.PE(data)
         except Exception, error:
@@ -156,7 +156,7 @@ class ContainerELF(Container):
         # Build executable instance
         try:
             if vm is not None:
-                self._executable = vm_load_elf(vm, filename)
+                self._executable = vm_load_elf(vm, data)
             else:
                 self._executable = elf_init.ELF(data)
         except Exception, error:
diff --git a/miasm2/analysis/sandbox.py b/miasm2/analysis/sandbox.py
index f7009f3e..ca6dcfe6 100644
--- a/miasm2/analysis/sandbox.py
+++ b/miasm2/analysis/sandbox.py
@@ -159,7 +159,7 @@ class OS_Win(OS):
                ]
 
     def __init__(self, custom_methods, *args, **kwargs):
-        from miasm2.jitter.loader.pe import vm_load_pe, preload_pe, libimp_pe
+        from miasm2.jitter.loader.pe import vm_load_pe, vm_load_pe_libs, preload_pe, libimp_pe
 
         super(OS_Win, self).__init__(custom_methods, *args, **kwargs)
 
@@ -173,19 +173,15 @@ class OS_Win(OS):
             all_pe = []
 
             # Load libs in memory
-            for dll_fname in self.ALL_IMP_DLL:
-                fname = os.path.join('win_dll', dll_fname)
-                e_lib = vm_load_pe(self.jitter.vm, fname)
-
-                libs.add_export_lib(e_lib, dll_fname)
-                all_pe.append(e_lib)
+            all_pe = vm_load_pe_libs(self.jitter.vm, self.ALL_IMP_DLL, libs)
 
             # Patch libs imports
-            for pe in all_pe:
+            for pe in all_pe.values():
                 preload_pe(self.jitter.vm, pe, libs)
 
         # Load main pe
-        self.pe = vm_load_pe(self.jitter.vm, self.fname)
+        with open(self.fname) as fstream:
+            self.pe = vm_load_pe(self.jitter.vm, fstream.read())
 
         # Fix pe imports
         preload_pe(self.jitter.vm, self.pe, libs)
@@ -225,17 +221,16 @@ class OS_Linux(OS):
         super(OS_Linux, self).__init__(custom_methods, *args, **kwargs)
 
         # Import manager
-        libs = libimp_elf()
-        self.libs = libs
+        self.libs = libimp_elf()
 
-        elf = vm_load_elf(self.jitter.vm, self.fname)
-        self.elf = elf
-        preload_elf(self.jitter.vm, elf, libs)
+        with open(self.fname) as fstream:
+            self.elf = vm_load_elf(self.jitter.vm, fstream.read())
+        preload_elf(self.jitter.vm, self.elf, self.libs)
 
-        self.entry_point = elf.Ehdr.entry
+        self.entry_point = self.elf.Ehdr.entry
 
         # Library calls handler
-        self.jitter.add_lib_handler(libs, custom_methods)
+        self.jitter.add_lib_handler(self.libs, custom_methods)
 
 class OS_Linux_str(OS):
     def __init__(self, custom_methods, *args, **kwargs):
diff --git a/miasm2/jitter/loader/elf.py b/miasm2/jitter/loader/elf.py
index 528ff4f5..3fd08725 100644
--- a/miasm2/jitter/loader/elf.py
+++ b/miasm2/jitter/loader/elf.py
@@ -47,13 +47,13 @@ def preload_elf(vm, e, runtime_lib, patch_vm_imp=True):
 
 
 
-def vm_load_elf(vm, fname, **kargs):
+def vm_load_elf(vm, fdata, **kargs):
     """
     Very dirty elf loader
     TODO XXX: implement real loader
     """
     #log.setLevel(logging.DEBUG)
-    e = elf_init.ELF(open(fname, 'rb').read(), **kargs)
+    e = elf_init.ELF(fdata, **kargs)
     i = interval()
     all_data = {}
     for p in e.ph.phlist:
diff --git a/miasm2/jitter/loader/pe.py b/miasm2/jitter/loader/pe.py
index f492cba3..6b19fc16 100644
--- a/miasm2/jitter/loader/pe.py
+++ b/miasm2/jitter/loader/pe.py
@@ -1,13 +1,15 @@
+import os
 import struct
+import logging
 from collections import defaultdict
 
 from elfesteem import pe
 from elfesteem import cstruct
 from elfesteem import *
+
 from miasm2.jitter.csts import *
-from utils import canon_libname_libfunc, libimp
+from miasm2.jitter.loader.utils import canon_libname_libfunc, libimp
 
-import logging
 
 log = logging.getLogger('loader_pe')
 hnd = logging.StreamHandler()
@@ -88,93 +90,128 @@ def get_export_name_addr_list(e):
 
 
 
-def vm_load_pe(vm, fname, align_s=True, load_hdr=True,
-               **kargs):
-    e = pe_init.PE(open(fname, 'rb').read(), **kargs)
+def vm_load_pe(vm, fdata, align_s=True, load_hdr=True, **kargs):
+    """Load a PE in memory (@vm) from a data buffer @fdata
+    @vm: VmMngr instance
+    @fdata: data buffer to parse
+    @align_s: (optional) If False, keep gaps between section
+    @load_hdr: (optional) If False, do not load the NThdr in memory
+    Return the corresponding PE instance.
 
+    Extra arguments are passed to PE instanciation.
+    If all sections are aligned, they will be mapped on several different pages
+    Otherwise, a big page is created, containing all sections
+    """
+    # Parse and build a PE instance
+    pe = pe_init.PE(fdata, **kargs)
+
+    # Check if all section are aligned
     aligned = True
-    for s in e.SHList:
-        if s.addr & 0xFFF:
+    for section in pe.SHList:
+        if section.addr & 0xFFF:
             aligned = False
             break
 
     if aligned:
+        # Loader NT header
         if load_hdr:
-            hdr_len = max(0x200, e.NThdr.sizeofheaders)
-            min_len = min(e.SHList[0].addr, 0x1000)#e.NThdr.sizeofheaders)
-            pe_hdr = e.content[:hdr_len]
-            pe_hdr = pe_hdr + min_len * "\x00"
-            pe_hdr = pe_hdr[:min_len]
-            vm.add_memory_page(
-                e.NThdr.ImageBase, PAGE_READ | PAGE_WRITE, pe_hdr)
+            # Header length
+            hdr_len = max(0x200, pe.NThdr.sizeofheaders)
+            # Page minimum size
+            min_len = min(pe.SHList[0].addr, 0x1000)
+
+            # Get and pad the pe_hdr
+            pe_hdr = pe.content[:hdr_len] + max(0, (min_len - hdr_len)) * "\x00"
+            vm.add_memory_page(pe.NThdr.ImageBase, PAGE_READ | PAGE_WRITE,
+                               pe_hdr)
+
+        # Align sections size
         if align_s:
-            for i, s in enumerate(e.SHList[:-1]):
-                s.size = e.SHList[i + 1].addr - s.addr
-                s.rawsize = s.size
-                s.data = strpatchwork.StrPatchwork(s.data[:s.size])
-                s.offset = s.addr
-            s = e.SHList[-1]
-            s.size = (s.size + 0xfff) & 0xfffff000
-        for s in e.SHList:
-            data = str(s.data)
-            data += "\x00" * (s.size - len(data))
-            # log.debug('SECTION %s %s' % (hex(s.addr),
-            # hex(e.rva2virt(s.addr))))
-            vm.add_memory_page(
-                e.rva2virt(s.addr), PAGE_READ | PAGE_WRITE, data)
-            # s.offset = s.addr
-        return e
-
-    # not aligned
-    log.warning('pe is not aligned, creating big section')
-    min_addr = None
+            # Use the next section address to compute the new size
+            for i, section in enumerate(pe.SHList[:-1]):
+                new_size = pe.SHList[i + 1].addr - section.addr
+                section.size = new_size
+                section.rawsize = new_size
+                section.data = strpatchwork.StrPatchwork(section.data[:new_size])
+                section.offset = section.addr
+
+            # Last section alignement
+            last_section = pe.SHList[-1]
+            last_section.size = (last_section.size + 0xfff) & 0xfffff000
+
+        # Pad sections with null bytes and map them
+        for section in pe.SHList:
+            data = str(section.data)
+            data += "\x00" * (section.size - len(data))
+            vm.add_memory_page(pe.rva2virt(section.addr),
+                               PAGE_READ | PAGE_WRITE, data)
+
+        return pe
+
+    # At least one section is not aligned
+    log.warning('PE is not aligned, creating big section')
+    min_addr = 0 if load_hdr else None
     max_addr = None
     data = ""
 
-    if load_hdr:
-        data = e.content[:0x400]
-        data += (e.SHList[0].addr - len(data)) * "\x00"
-        min_addr = 0
-
-    for i, s in enumerate(e.SHList):
-        if i < len(e.SHList) - 1:
-            s.size = e.SHList[i + 1].addr - s.addr
-        s.rawsize = s.size
-        s.offset = s.addr
-
-        if min_addr is None or s.addr < min_addr:
-            min_addr = s.addr
-        if max_addr is None or s.addr + s.size > max_addr:
-            max_addr = s.addr + max(s.size, len(s.data))
-    min_addr = e.rva2virt(min_addr)
-    max_addr = e.rva2virt(max_addr)
-    log.debug('%s %s %s' %
-              (hex(min_addr), hex(max_addr), hex(max_addr - min_addr)))
-
+    for i, section in enumerate(pe.SHList):
+        if i < len(pe.SHList) - 1:
+            # If it is not the last section, use next section address
+            section.size = pe.SHList[i + 1].addr - section.addr
+        section.rawsize = section.size
+        section.offset = section.addr
+
+        # 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))
+
+    min_addr = pe.rva2virt(min_addr)
+    max_addr = pe.rva2virt(max_addr)
+    log.debug('Min: 0x%x, Max: 0x%x, Size: 0x%x' % (min_addr, max_addr,
+                                                    (max_addr - min_addr)))
+
+    # Create only one big section containing the whole PE
     vm.add_memory_page(min_addr,
-                          PAGE_READ | PAGE_WRITE,
-                          (max_addr - min_addr) * "\x00")
-    for s in e.SHList:
-        log.debug('%s %s' % (hex(e.rva2virt(s.addr)), len(s.data)))
-        vm.set_mem(e.rva2virt(s.addr), str(s.data))
-    return e
-
-
-def vm_load_pe_lib(fname_in, libs, lib_path_base, patch_vm_imp, **kargs):
+                       PAGE_READ | PAGE_WRITE,
+                       (max_addr - min_addr) * "\x00")
+
+    # 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))
+
+    return pe
+
+
+def vm_load_pe_lib(vm, fname_in, libs, lib_path_base, **kargs):
+    """Call vm_load_pe on @fname_in and update @libs accordingly
+    @vm: VmMngr instance
+    @fname_in: library name
+    @libs: libimp_pe instance
+    @lib_path_base: DLLs relative path
+    Return the corresponding PE instance
+    Extra arguments are passed to vm_load_pe
+    """
     fname = os.path.join(lib_path_base, fname_in)
-    e = vm_load_pe(fname, **kargs)
-    libs.add_export_lib(e, fname_in)
-    # preload_pe(e, libs, patch_vm_imp)
-    return e
-
-
-def vm_load_pe_libs(libs_name, libs, lib_path_base="win_dll",
-                    patch_vm_imp=True, **kargs):
-    lib_imgs = {}
-    for fname in libs_name:
-        e = vm_load_pe_lib(fname, libs, lib_path_base, patch_vm_imp)
-        lib_imgs[fname] = e
-    return lib_imgs
+    with open(fname) as fstream:
+        pe = vm_load_pe(vm, fstream.read(), **kargs)
+    libs.add_export_lib(pe, fname_in)
+    return pe
+
+
+def vm_load_pe_libs(vm, libs_name, libs, lib_path_base="win_dll", **kargs):
+    """Call vm_load_pe_lib on each @libs_name filename
+    @vm: VmMngr instance
+    @libs_name: list of str
+    @libs: libimp_pe instance
+    @lib_path_base: (optional) DLLs relative path
+    Return a dictionnary Filename -> PE instances
+    Extra arguments are passed to vm_load_pe_lib
+    """
+    return {fname: vm_load_pe_lib(vm, fname, libs, lib_path_base, **kargs)
+            for fname in libs_name}
 
 
 def vm_fix_imports_pe_libs(lib_imgs, libs, lib_path_base="win_dll",