Merge pull request #1222 from serpilliere/fix_str_api

Fix str api

2 files changed, 11 insertions, 11 deletions

diff --git a/miasm/os_dep/common.py b/miasm/os_dep/common.py
index 4a92ef2a..74100817 100644
--- a/miasm/os_dep/common.py
+++ b/miasm/os_dep/common.py
@@ -71,15 +71,15 @@ class heap(object):
         self.addr &= self.mask ^ (self.align - 1)
         return ret
 
-    def alloc(self, jitter, size, perm=PAGE_READ | PAGE_WRITE):
+    def alloc(self, jitter, size, perm=PAGE_READ | PAGE_WRITE, cmt=""):
         """
         @jitter: a jitter instance
         @size: the size to allocate
         @perm: permission flags (see vm_alloc doc)
         """
-        return self.vm_alloc(jitter.vm, size, perm)
+        return self.vm_alloc(jitter.vm, size, perm=perm, cmt=cmt)
 
-    def vm_alloc(self, vm, size, perm=PAGE_READ | PAGE_WRITE):
+    def vm_alloc(self, vm, size, perm=PAGE_READ | PAGE_WRITE, cmt=""):
         """
         @vm: a VmMngr instance
         @size: the size to allocate
@@ -91,7 +91,7 @@ class heap(object):
             addr,
             perm,
             b"\x00" * (size),
-            "Heap alloc by %s" % get_caller_name(2)
+            "Heap alloc by %s %s" % (get_caller_name(2), cmt)
         )
         return addr
 
diff --git a/miasm/os_dep/win_api_x86_32.py b/miasm/os_dep/win_api_x86_32.py
index e42f8006..9bf3975b 100644
--- a/miasm/os_dep/win_api_x86_32.py
+++ b/miasm/os_dep/win_api_x86_32.py
@@ -272,7 +272,7 @@ class mdl(object):
 
 def kernel32_HeapAlloc(jitter):
     ret_ad, args = jitter.func_args_stdcall(["heap", "flags", "size"])
-    alloc_addr = winobjs.heap.alloc(jitter, args.size)
+    alloc_addr = winobjs.heap.alloc(jitter, args.size, cmt=hex(ret_ad))
     jitter.func_ret_stdcall(ret_ad, alloc_addr)
 
 
@@ -1120,7 +1120,7 @@ def shell32_CommandLineToArgvW(jitter):
     addr_ret = winobjs.heap.alloc(jitter, 4 * (len(tks) + 1))
     o = 0
     for i, t in enumerate(tks):
-        jitter.set_win_str_w(addr + o, t)
+        set_win_str_w(jitter, addr + o, t)
         jitter.vm.set_u32(addr_ret + 4 * i, addr + o)
         o += len(t)*2 + 2
 
@@ -1335,7 +1335,7 @@ def ntoskrnl_RtlGetVersion(jitter):
                     0x2,  # min vers
                     0x666,  # build nbr
                     0x2,   # platform id
-                    ) + jitter.set_win_str_w("Service pack 4")
+                    ) + encode_win_str_w("Service pack 4")
 
     jitter.vm.set_mem(args.ptr_version, s)
     jitter.func_ret_stdcall(ret_ad, 0)
@@ -1521,7 +1521,7 @@ def kernel32_lstrcpy(jitter):
 def msvcrt__mbscpy(jitter):
     ret_ad, args = jitter.func_args_cdecl(["ptr_str1", "ptr_str2"])
     s2 = get_win_str_w(jitter, args.ptr_str2)
-    jitter.set_win_str_w(args.ptr_str1, s2)
+    set_win_str_w(jitter, args.ptr_str1, s2)
     jitter.func_ret_cdecl(ret_ad, args.ptr_str1)
 
 def msvcrt_wcscpy(jitter):
@@ -1535,7 +1535,7 @@ def kernel32_lstrcpyn(jitter):
     if len(s2) >= args.mlen:
         s2 = s2[:args.mlen - 1]
     log.info("Copy '%r'", s2)
-    jitter.set_win_str_a(args.ptr_str1, s2)
+    set_win_str_a(jitter, args.ptr_str1, s2)
     jitter.func_ret_stdcall(ret_ad, args.ptr_str1)
 
 
@@ -2889,10 +2889,10 @@ class win32_find_data(object):
                         self.filesizelow,
                         self.dwreserved0,
                         self.dwreserved1)
-        fname = self.cfilename.encode('utf-8') + b'\x00' * MAX_PATH
+        fname = encode_win_str_w(self.cfilename) + b'\x00' * MAX_PATH
         fname = fname[:MAX_PATH]
         s += fname
-        fname = self.alternamefilename.encode('utf-8') + b'\x00' * 14
+        fname = encode_win_str_w(self.alternamefilename) + b'\x00' * 14
         fname = fname[:14]
         s += fname
         return s