Support python2/python3

1 files changed, 51 insertions, 31 deletions

diff --git a/miasm2/os_dep/common.py b/miasm2/os_dep/common.py
index 7e46b276..ed68185f 100644
--- a/miasm2/os_dep/common.py
+++ b/miasm2/os_dep/common.py
@@ -1,5 +1,8 @@
 import os
 
+from future.utils import viewitems
+
+from miasm2.core.utils import force_bytes
 from miasm2.jitter.csts import PAGE_READ, PAGE_WRITE
 from miasm2.core.utils import get_caller_name
 from miasm2.core.utils import pck64, upck64
@@ -11,7 +14,7 @@ def get_str_ansi(jitter, ad_str, max_char=None):
     l = 0
     tmp = ad_str
     while ((max_char is None or l < max_char) and
-           jitter.vm.get_mem(tmp, 1) != "\x00"):
+           jitter.vm.get_mem(tmp, 1) != b"\x00"):
         tmp += 1
         l += 1
     return jitter.vm.get_mem(ad_str, l)
@@ -21,22 +24,25 @@ def get_str_unic(jitter, ad_str, max_char=None):
     l = 0
     tmp = ad_str
     while ((max_char is None or l < max_char) and
-           jitter.vm.get_mem(tmp, 2) != "\x00\x00"):
+           jitter.vm.get_mem(tmp, 2) != b"\x00\x00"):
         tmp += 2
         l += 2
     s = jitter.vm.get_mem(ad_str, l)
-    # TODO: real unicode decoding
-    s = s[::2]
+    s = s.decode("utf-16le")
     return s
 
 
-def set_str_ansi(s):
-    return s + "\x00"
+def set_str_ansi(value):
+    value = force_bytes(value)
+    return value + b"\x00"
 
 
-def set_str_unic(s):
-    # TODO: real unicode encoding
-    return "\x00".join(list(s)) + '\x00' * 3
+def set_str_unic(value):
+    try:
+        value = value.decode()
+    except AttributeError:
+        pass
+    return value.encode("utf-16le") + b'\x00' * 2
 
 
 class heap(object):
@@ -74,30 +80,34 @@ class heap(object):
             combination of them); default is PAGE_READ|PAGE_WRITE
         """
         addr = self.next_addr(size)
-        vm.add_memory_page(addr, perm, "\x00" * (size),
-                           "Heap alloc by %s" % get_caller_name(2))
+        vm.add_memory_page(
+            addr,
+            perm,
+            b"\x00" * (size),
+            "Heap alloc by %s" % get_caller_name(2)
+        )
         return addr
 
     def get_size(self, vm, ptr):
         """
         @vm: a VmMngr instance
         @size: ptr to get the size of the associated allocation.
-        
+
         `ptr` can be the base address of a previous allocation, or an address
         within the allocated range. The size of the whole allocation is always
         returned, regardless ptr is the base address or not.
         """
-	assert vm.is_mapped(ptr, 1)
-	data = vm.get_all_memory()
-	ptr_page = data.get(ptr, None)
-	if ptr_page is None:
-	    for address, page_info in data.iteritems():
-		if address <= ptr < address + page_info["size"]:
-		    ptr_page = page_info
-		    break
-	    else:
-		raise RuntimeError("Must never happen (unmapped but mark as mapped by API)")
-	return ptr_page["size"]
+        assert vm.is_mapped(ptr, 1)
+        data = vm.get_all_memory()
+        ptr_page = data.get(ptr, None)
+        if ptr_page is None:
+            for address, page_info in viewitems(data):
+                if address <= ptr < address + page_info["size"]:
+                    ptr_page = page_info
+                    break
+            else:
+                raise RuntimeError("Must never happen (unmapped but mark as mapped by API)")
+        return ptr_page["size"]
 
 
 def windows_to_sbpath(path):
@@ -118,26 +128,36 @@ def unix_to_sbpath(path):
     return os.path.join(BASE_SB_PATH, *path)
 
 def get_fmt_args(fmt, cur_arg, get_str, get_arg_n):
-    output = ""
     idx = 0
     fmt = get_str(fmt)
+    if isinstance(fmt, bytes):
+        chars_format = b'%cdfsuxX'
+        char_percent = b'%'
+        char_string = b's'
+        output = b""
+    else:
+        chars_format = u'%cdfsuxX'
+        char_percent = u'%'
+        char_string = u's'
+        output = u""
+
     while True:
         if idx == len(fmt):
             break
-        char = fmt[idx]
+        char = fmt[idx:idx+1]
         idx += 1
-        if char == '%':
-            token = '%'
+        if char == char_percent:
+            token = char_percent
             while True:
-                char = fmt[idx]
+                char = fmt[idx:idx+1]
                 idx += 1
                 token += char
-                if char.lower() in '%cdfsux':
+                if char in chars_format:
                     break
-            if char == '%':
+            if char == char_percent:
                 output += char
                 continue
-            if token.endswith('s'):
+            if token.endswith(char_string):
                 addr = get_arg_n(cur_arg)
                 arg = get_str(addr)
             else: