Python2/3 String Compatibility Fix

9 files changed, 40 insertions, 24 deletions

diff --git a/python/__init__.py b/python/__init__.py
index 36ed4f0d..9d183465 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -51,6 +51,20 @@ def with_metaclass(meta, *bases):
     return type.__new__(metaclass, 'temporary_class', (), {})
 
 
+def cstr(arg):
+	if isinstance(arg, bytes) or arg is None:
+		return arg
+	else:
+		return arg.encode('charmap')
+
+
+def pyNativeStr(arg):
+	if isinstance(arg, str):
+		return arg
+	else:
+		return arg.decode('charmap')
+
+
 # Binary Ninja components
 import binaryninja._binaryninjacore as core
 # __all__ = [
diff --git a/python/binaryview.py b/python/binaryview.py
index c4fddd7c..25953b34 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -98,7 +98,7 @@ class StringReference(object):
 
 	@property
 	def value(self):
-		return self.view.read(self.start, self.length).decode("charmap")
+		return binaryninja.pyNativeStr(self.view.read(self.start, self.length))
 
 	def __repr__(self):
 		return "<%s: %#x, len %#x>" % (self.type, self.start, self.length)
@@ -3503,7 +3503,7 @@ class BinaryView(object):
 	def get_unique_section_names(self, name_list):
 		incoming_names = (ctypes.c_char_p * len(name_list))()
 		for i in range(0, len(name_list)):
-			incoming_names[i] = name_list[i].encode('charmap')
+			incoming_names[i] = binaryninja.cstr(name_list[i])
 		outgoing_names = core.BNGetUniqueSectionNames(self.handle, incoming_names, len(name_list))
 		result = []
 		for i in range(0, len(name_list)):
diff --git a/python/databuffer.py b/python/databuffer.py
index 298d0002..2fb5382d 100644
--- a/python/databuffer.py
+++ b/python/databuffer.py
@@ -23,6 +23,9 @@ import ctypes
 # Binary Ninja components
 from binaryninja import _binaryninjacore as core
 
+# 2-3 compatibility
+from binaryninja import pyNativeStr
+
 
 class DataBuffer(object):
 	def __init__(self, contents="", handle=None):
@@ -114,10 +117,7 @@ class DataBuffer(object):
 	def __str__(self):
 		buf = ctypes.create_string_buffer(len(self))
 		ctypes.memmove(buf, core.BNGetDataBufferContents(self.handle), len(self))
-		if isinstance(buf.raw, str):
-			return buf.raw
-		else:
-			return buf.raw.decode("charmap")
+		return pyNativeStr(buf.raw)
 
 	def __bytes__(self):
 		buf = ctypes.create_string_buffer(len(self))
diff --git a/python/demangle.py b/python/demangle.py
index 466324bc..2fa1028f 100644
--- a/python/demangle.py
+++ b/python/demangle.py
@@ -22,9 +22,11 @@ import ctypes
 
 # Binary Ninja components
 from binaryninja import _binaryninjacore as core
+from binaryninja import types
 
 # 2-3 compatibility
 from binaryninja import range
+from binaryninja import pyNativeStr
 
 
 def get_qualified_name(names):
@@ -58,28 +60,26 @@ def demangle_ms(arch, mangled_name):
 		(<type: public: static enum Foobar::foo __cdecl (enum Foobar::foo)>, ['Foobar', 'testf'])
 		>>>
 	"""
-	from binaryninja import types
 	handle = ctypes.POINTER(core.BNType)()
 	outName = ctypes.POINTER(ctypes.c_char_p)()
 	outSize = ctypes.c_ulonglong()
 	names = []
 	if core.BNDemangleMS(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize)):
 		for i in range(outSize.value):
-			names.append(outName[i].decode('charmap'))
+			names.append(pyNativeStr(outName[i]))
 		core.BNFreeDemangledName(ctypes.byref(outName), outSize.value)
 		return (types.Type(handle), names)
 	return (None, mangled_name)
 
 
 def demangle_gnu3(arch, mangled_name):
-	from binaryninja import types
 	handle = ctypes.POINTER(core.BNType)()
 	outName = ctypes.POINTER(ctypes.c_char_p)()
 	outSize = ctypes.c_ulonglong()
 	names = []
 	if core.BNDemangleGNU3(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize)):
 		for i in range(outSize.value):
-			names.append(outName[i].decode('charmap'))
+			names.append(pyNativeStr(outName[i]))
 		core.BNFreeDemangledName(ctypes.byref(outName), outSize.value)
 		if not handle:
 			return (None, names)
diff --git a/python/downloadprovider.py b/python/downloadprovider.py
index 09abd516..14368ee7 100644
--- a/python/downloadprovider.py
+++ b/python/downloadprovider.py
@@ -31,6 +31,9 @@ from binaryninja import with_metaclass
 from binaryninja import startup
 from binaryninja import log
 
+# 2-3 compatibility
+from binaryninja import pyNativeStr
+
 
 class DownloadInstance(object):
 	def __init__(self, provider, handle = None):
@@ -166,7 +169,7 @@ if sys.version_info >= (2, 7, 9):
 					opener = build_opener(ProxyHandler({'https': proxy_setting}))
 					install_opener(opener)
 
-				r = urlopen(url.decode("charmap"))
+				r = urlopen(pyNativeStr(url))
 				total_size = int(r.headers.get('content-length', 0))
 				bytes_sent = 0
 				while True:
@@ -224,7 +227,7 @@ else:
 					else:
 						proxies = None
 
-					r = requests.get(url.decode("charmap"), proxies=proxies)
+					r = requests.get(pyNativeStr(url), proxies=proxies)
 					if not r.ok:
 						core.BNSetErrorForDownloadInstance(self.handle, "Received error from server")
 						return -1
diff --git a/python/generator.cpp b/python/generator.cpp
index 2da856ff..afc77254 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -198,13 +198,9 @@ int main(int argc, char* argv[])
 	fprintf(out, "\t_base_path = os.path.join(os.path.dirname(__file__), \"..\", \"..\")\n");
 	fprintf(out, "\tcore = ctypes.CDLL(os.path.join(_base_path, \"binaryninjacore.dll\"))\n");
 	fprintf(out, "else:\n");
-	fprintf(out, "\traise Exception(\"OS not supported\")\n\n");
+	fprintf(out, "\traise Exception(\"OS not supported\")\n\n\n");
 
-	fprintf(out, "def cstr(arg):\n");
-	fprintf(out, "\tif isinstance(arg, str):\n");
-	fprintf(out, "\t\treturn arg.encode('charmap')\n");
-	fprintf(out, "\telse:\n");
-	fprintf(out, "\t\treturn arg\n\n");
+	fprintf(out, "from binaryninja import cstr, pyNativeStr\n\n\n");
 
 	// Create type objects
 	fprintf(out, "# Type definitions\n");
@@ -226,8 +222,8 @@ int main(int argc, char* argv[])
 					(arg.type->GetChildType()->GetWidth() == 1) &&
 					(arg.type->GetChildType()->IsSigned()))
 					{
-						fprintf(out, "\t@property\n\tdef %s(self):\n\t\treturn self._%s.decode('charmap')\n", arg.name.c_str(), arg.name.c_str());
-						fprintf(out, "\t@%s.setter\n\tdef %s(self, value):\n\t\tself._%s = value.encode('charmap')\n", arg.name.c_str(), arg.name.c_str(), arg.name.c_str());
+						fprintf(out, "\t@property\n\tdef %s(self):\n\t\treturn pyNativeStr(self._%s)\n", arg.name.c_str(), arg.name.c_str());
+						fprintf(out, "\t@%s.setter\n\tdef %s(self, value):\n\t\tself._%s = cstr(value)\n", arg.name.c_str(), arg.name.c_str(), arg.name.c_str());
 						stringField = true;
 					}
 			}
@@ -423,7 +419,7 @@ int main(int argc, char* argv[])
 			}
 			else
 				fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
-			fprintf(out, "\tstring = str(ctypes.cast(result, ctypes.c_char_p).value.decode('charmap'))\n");
+			fprintf(out, "\tstring = str(pyNativeStr(ctypes.cast(result, ctypes.c_char_p).value))\n");
 			fprintf(out, "\tBNFreeString(result)\n");
 			fprintf(out, "\treturn string\n");
 		}
diff --git a/python/metadata.py b/python/metadata.py
index b17bbf4f..606e4ceb 100644
--- a/python/metadata.py
+++ b/python/metadata.py
@@ -28,6 +28,7 @@ from binaryninja.enums import MetadataType
 
 # 2-3 compatibility
 from binaryninja import range
+from binaryninja import pyNativeStr
 
 
 class Metadata(object):
@@ -148,7 +149,7 @@ class Metadata(object):
 			try:
 				for i in range(result.contents.size):
 					if isinstance(result.contents.keys[i], bytes):
-						yield str(result.contents.keys[i].decode('charmap'))
+						yield str(pyNativeStr(result.contents.keys[i]))
 					else:
 						yield result.contents.keys[i]
 			finally:
diff --git a/python/setting.py b/python/setting.py
index 355492aa..93f6b72a 100644
--- a/python/setting.py
+++ b/python/setting.py
@@ -25,6 +25,7 @@ from binaryninja import _binaryninjacore as core
 
 # 2-3 compatibility
 from binaryninja import range
+from binaryninja import pyNativeStr
 
 
 class Setting(object):
@@ -62,7 +63,7 @@ class Setting(object):
 		result = core.BNSettingGetStringList(self.plugin_name, name, default_list, ctypes.byref(length))
 		out_list = []
 		for i in range(length.value):
-			out_list.append(result[i].decode('charmap'))
+			out_list.append(pyNativeStr(result[i]))
 		core.BNFreeStringList(result, length)
 		return out_list
 
diff --git a/python/types.py b/python/types.py
index e8ce7464..12e7733f 100644
--- a/python/types.py
+++ b/python/types.py
@@ -30,6 +30,7 @@ from binaryninja.enums import SymbolType, TypeClass, NamedTypeReferenceClass, In
 
 # 2-3 compatibility
 from binaryninja import range
+from binaryninja import pyNativeStr
 
 
 class QualifiedName(object):
@@ -41,7 +42,7 @@ class QualifiedName(object):
 			self.name = name.name
 			self.byte_name = [n.encode('charmap') for n in name.name]
 		else:
-			self.name = [i.decode('charmap') for i in name]
+			self.name = [pyNativeStr(i) for i in name]
 			self.byte_name = name
 
 	def __str__(self):