Various Python 3 support changes

1 files changed, 134 insertions, 45 deletions

diff --git a/python/generator.cpp b/python/generator.cpp
index 8fc415b0..2da856ff 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -121,10 +121,8 @@ void OutputType(FILE* out, Type* type, bool isReturnType = false, bool isCallbac
 		{
 			if (isReturnType)
 				fprintf(out, "ctypes.POINTER(ctypes.c_byte)");
-			else {
-				//fprintf(out, "compatstring");
+			else
 				fprintf(out, "ctypes.c_char_p");
-			}
 			break;
 		}
 		else if (type->GetChildType()->GetClass() == FunctionTypeClass)
@@ -175,7 +173,7 @@ int main(int argc, char* argv[])
 	}
 
 	bool ok = arch->GetStandalonePlatform()->ParseTypesFromSourceFile(argv[1], types, vars, funcs, errors);
-	fprintf(stderr, "Errors: %s", errors.c_str());
+	fprintf(stderr, "Errors: %s\n", errors.c_str());
 	if (!ok)
 		return 1;
 
@@ -190,7 +188,6 @@ int main(int argc, char* argv[])
 	fprintf(out, "import platform\n");
 	fprintf(out, "core = None\n");
 	fprintf(out, "_base_path = None\n");
-	fprintf(out, "ctypes.set_conversion_mode('utf-8', 'strict')\n");
 	fprintf(out, "if platform.system() == \"Darwin\":\n");
 	fprintf(out, "\t_base_path = os.path.join(os.path.dirname(__file__), \"..\", \"..\", \"..\", \"MacOS\")\n");
 	fprintf(out, "\tcore = ctypes.CDLL(os.path.join(_base_path, \"libbinaryninjacore.dylib\"))\n\n");
@@ -202,40 +199,12 @@ int main(int argc, char* argv[])
 	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, "class compatstring(ctypes.c_char_p):\n");
-	// fprintf(out, "\tdef __init__(self, value=None):\n");
-	// fprintf(out, "\t\tsuper(compatstring, self).__init__()\n");
-	// fprintf(out, "\t\tif value is not None:\n");
-	// fprintf(out, "\t\t\tself.value = value\n");
-	// fprintf(out, "\t@classmethod\n");
-	// fprintf(out, "\tdef from_param(cls, value):\n");
-	// fprintf(out, "\t\tif not isinstance(value, bytes):\n");
-	// fprintf(out, "\t\t\treturn super(compatstring, cls).from_param(value.encode('utf8'))\n");
-	// fprintf(out, "\t\treturn super(compatstring, cls).from_param(value)\n\n");
-	// fprintf(out, "\t@property\n");
-	// fprintf(out, "\tdef value(self, value):\n");
-	// fprintf(out, "\t\tif not isinstance(value, bytes):\n");
-	// fprintf(out, "\t\t\treturn super(compatstring, cls).from_param(value.encode('utf8'))\n");
-	// fprintf(out, "\t\treturn super(compatstring, cls).from_param(value)\n\n");
-	fprintf(out, "class compatstring(ctypes.c_char_p):\n");
-	fprintf(out, "	@classmethod\n");
-	fprintf(out, "	def from_param(cls, obj):\n");
-	fprintf(out, "		if (obj is not None) and (not isinstance(obj, cls)):\n");
-	fprintf(out, "			if not isinstance(obj, basestring):\n");
-	fprintf(out, "				raise TypeError('parameter must be a string type instance')\n");
-	fprintf(out, "			if not isinstance(obj, unicode):\n");
-	fprintf(out, "				obj = unicode(obj)\n");
-	fprintf(out, "			obj = obj.encode('utf-8')\n");
-	fprintf(out, "		return ctypes.c_char_p.from_param(obj)\n");
-	fprintf(out, "\n");
-	fprintf(out, "	def decode(self):\n");
-	fprintf(out, "		if self.value is None:\n");
-	fprintf(out, "			return None\n");
-	fprintf(out, "		return self.value.decode('utf-8')\n");
-	fprintf(out, "	@property\n");
-	fprintf(out, "	def value(self, c_void_p=ctypes.c_void_p):\n");
-	fprintf(out, "		addr = c_void_p.from_buffer(self).value\n");
-	fprintf(out, "		return \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");
 
 	// Create type objects
 	fprintf(out, "# Type definitions\n");
@@ -248,7 +217,23 @@ int main(int argc, char* argv[])
 		if (i.second->GetClass() == StructureTypeClass)
 		{
 			fprintf(out, "class %s(ctypes.Structure):\n", name.c_str());
-			fprintf(out, "\tpass\n");
+
+			// python uses str's, C uses byte-arrays
+			bool stringField = false;
+			for (auto& arg : i.second->GetStructure()->GetMembers())
+			{
+				if ((arg.type->GetClass() == PointerTypeClass) &&
+					(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());
+						stringField = true;
+					}
+			}
+
+			if (!stringField)
+				fprintf(out, "\tpass\n");
 		}
 		else if (i.second->GetClass() == EnumerationTypeClass)
 		{
@@ -313,7 +298,15 @@ int main(int argc, char* argv[])
 				fprintf(out, "%s._fields_ = [\n", name.c_str());
 				for (auto& j : type->GetStructure()->GetMembers())
 				{
-					fprintf(out, "\t\t(\"%s\", ", j.name.c_str());
+					// To help the python->C wrappers
+					if ((j.type->GetClass() == PointerTypeClass) &&
+						(j.type->GetChildType()->GetWidth() == 1) &&
+						(j.type->GetChildType()->IsSigned()))
+						{
+							fprintf(out, "\t\t(\"_%s\", ", j.name.c_str());
+						}
+					else
+						fprintf(out, "\t\t(\"%s\", ", j.name.c_str());
 					OutputType(out, j.type);
 					fprintf(out, "),\n");
 				}
@@ -347,6 +340,22 @@ int main(int argc, char* argv[])
 			(i.second->GetChildType()->GetChildType()->IsSigned());
 		// Pointer returns will be automatically wrapped to return None on null pointer
 		bool pointerResult = (i.second->GetChildType()->GetClass() == PointerTypeClass);
+
+		// From python -> C python3 requires str -> str.encode('charmap')
+		bool stringArgument = false;
+		for (auto& arg : i.second->GetParameters())
+		{
+			if ((arg.type->GetClass() == PointerTypeClass) &&
+				(arg.type->GetChildType()->GetWidth() == 1) &&
+				(arg.type->GetChildType()->IsSigned()))
+				{
+					stringArgument = true;
+					break;
+				}
+		}
+		if (name == "BNFreeString")
+			stringArgument = false;
+
 		bool callbackConvention = false;
 		if (name == "BNAllocString")
 		{
@@ -358,7 +367,7 @@ int main(int argc, char* argv[])
 		}
 
 		string funcName = name;
-		if (stringResult || pointerResult)
+		if (stringResult || pointerResult || stringArgument)
 			funcName = string("_") + funcName;
 
 		fprintf(out, "%s = core.%s\n", funcName.c_str(), name.c_str());
@@ -391,8 +400,30 @@ int main(int argc, char* argv[])
 		{
 			// Emit wrapper to get Python string and free native memory
 			fprintf(out, "def %s(*args):\n", name.c_str());
-			fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
-			fprintf(out, "\tstring = ctypes.cast(result, ctypes.c_char_p).value\n");
+			if (stringArgument)
+			{
+				fprintf(out, "\tresult = %s(", funcName.c_str());
+				string stringArgFuncCall = "";
+				size_t argN = 0;
+				for (auto& arg : i.second->GetParameters())
+				{
+					if ((arg.type->GetClass() == PointerTypeClass) &&
+						(arg.type->GetChildType()->GetWidth() == 1) &&
+						(arg.type->GetChildType()->IsSigned()))
+						{
+							stringArgFuncCall += "cstr(args[" + to_string(argN) + "]), ";
+						}
+					else
+						stringArgFuncCall += "args[" + to_string(argN) + "], ";
+					argN++;
+				}
+				stringArgFuncCall = stringArgFuncCall.substr(0, stringArgFuncCall.size()-2);
+				stringArgFuncCall += ")\n";
+				fprintf(out, "%s", stringArgFuncCall.c_str());
+			}
+			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, "\tBNFreeString(result)\n");
 			fprintf(out, "\treturn string\n");
 		}
@@ -400,11 +431,58 @@ int main(int argc, char* argv[])
 		{
 			// Emit wrapper to return None on null pointer
 			fprintf(out, "def %s(*args):\n", name.c_str());
-			fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
+			if (stringArgument)
+			{
+				fprintf(out, "\tresult = %s(", funcName.c_str());
+				string stringArgFuncCall = "";
+				size_t argN = 0;
+				for (auto& arg : i.second->GetParameters())
+				{
+					if ((arg.type->GetClass() == PointerTypeClass) &&
+						(arg.type->GetChildType()->GetWidth() == 1) &&
+						(arg.type->GetChildType()->IsSigned()))
+						{
+							stringArgFuncCall += "cstr(args[" + to_string(argN) + "]), ";
+						}
+					else
+						stringArgFuncCall += "args[" + to_string(argN) + "], ";
+					argN++;
+				}
+				stringArgFuncCall = stringArgFuncCall.substr(0, stringArgFuncCall.size()-2);
+				stringArgFuncCall += ")\n";
+				fprintf(out, "%s", stringArgFuncCall.c_str());
+			}
+			else
+				fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
 			fprintf(out, "\tif not result:\n");
 			fprintf(out, "\t\treturn None\n");
 			fprintf(out, "\treturn result\n");
 		}
+		else if (stringArgument)
+		{
+			fprintf(out, "def %s(*args):\n", name.c_str());
+			{
+				fprintf(out, "\treturn %s(", funcName.c_str());
+				string stringArgFuncCall = "";
+				size_t argN = 0;
+				for (auto& arg : i.second->GetParameters())
+				{
+					if ((arg.type->GetClass() == PointerTypeClass) &&
+						(arg.type->GetChildType()->GetWidth() == 1) &&
+						(arg.type->GetChildType()->IsSigned()))
+						{
+							stringArgFuncCall += "cstr(args[" + to_string(argN) + "]), ";
+						}
+					else
+						stringArgFuncCall += "args[" + to_string(argN) + "], ";
+					argN++;
+				}
+				stringArgFuncCall = stringArgFuncCall.substr(0, stringArgFuncCall.size()-2);
+				stringArgFuncCall += ")\n";
+				fprintf(out, "%s", stringArgFuncCall.c_str());
+			}
+		}
+		fprintf(out, "\n");
 	}
 
 	fprintf(out, "\n# Helper functions\n");
@@ -413,6 +491,17 @@ int main(int argc, char* argv[])
 	fprintf(out, "\t\treturn ctypes.cast(value, ctypes.POINTER(handle_type))\n");
 	fprintf(out, "\traise ValueError('expected pointer to %%s' %% str(handle_type))\n");
 
+	// The following method is addapted from python/enum/__init__.py, lines 25-36
+	fprintf(out, "\ntry:\n");
+	fprintf(out, "\tbasestring\n");
+	fprintf(out, "\tunicode\n");
+	fprintf(out, "except NameError:\n");
+	fprintf(out, "\t# In Python 2 basestring is the ancestor of both str and unicode\n");
+	fprintf(out, "\t# in Python 3 it's just str, but was missing in 3.1\n");
+	fprintf(out, "\t# In Python 3 unicode no longer exists (it's just str)\n");
+	fprintf(out, "\tbasestring = str\n");
+	fprintf(out, "\tunicode = str\n");
+
 	fprintf(out, "\n# Set path for core plugins\n");
 	fprintf(out, "BNSetBundledPluginDirectory(os.path.join(_base_path, \"plugins\"))\n");