1 files changed, 56 insertions, 49 deletions

diff --git a/python/generator.cpp b/python/generator.cpp
index 485655c2..1c0e7b16 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -98,7 +98,17 @@ void OutputType(FILE* out, Type* type, bool isReturnType = false, bool isCallbac
 			fprintf(out, "ctypes.c_double");
 		break;
 	case NamedTypeReferenceClass:
-		fprintf(out, "%s", type->GetNamedTypeReference()->GetName().GetString().c_str());
+		if (type->GetNamedTypeReference()->GetTypeClass() == EnumNamedTypeClass)
+		{
+			string name = type->GetNamedTypeReference()->GetName().GetString();
+			if (name.size() > 2 && name.substr(0, 2) == "BN")
+				name = name.substr(2);
+			fprintf(out, "%sEnum", name.c_str());
+		}
+		else
+		{
+			fprintf(out, "%s", type->GetNamedTypeReference()->GetName().GetString().c_str());
+		}
 		break;
 	case PointerTypeClass:
 		if (isCallback || (type->GetChildType()->GetClass() == VoidTypeClass))
@@ -144,9 +154,9 @@ void OutputType(FILE* out, Type* type, bool isReturnType = false, bool isCallbac
 
 int main(int argc, char* argv[])
 {
-	if (argc < 3)
+	if (argc < 4)
 	{
-		fprintf(stderr, "Usage: generator <header> <output>\n");
+		fprintf(stderr, "Usage: generator <header> <output> <output_enum>\n");
 		return 1;
 	}
 
@@ -161,27 +171,30 @@ int main(int argc, char* argv[])
 		return 1;
 
 	FILE* out = fopen(argv[2], "w");
+	FILE* enums = fopen(argv[3], "w");
 
+	fprintf(out, "from __future__ import absolute_import\n");
 	fprintf(out, "import ctypes, os\n\n");
+	fprintf(enums, "import enum");
 
 	fprintf(out, "# Load core module\n");
-#if defined(__APPLE__)
-	fprintf(out, "_base_path = os.path.join(os.path.dirname(__file__), \"..\", \"..\", \"..\", \"MacOS\")\n");
-#else
-	fprintf(out, "_base_path = os.path.join(os.path.dirname(__file__), \"..\", \"..\")\n");
-#endif
-
-#ifdef WIN32
-	fprintf(out, "core = ctypes.CDLL(os.path.join(_base_path, \"binaryninjacore.dll\"))\n\n");
-#elif defined(__APPLE__)
-	fprintf(out, "core = ctypes.CDLL(os.path.join(_base_path, \"libbinaryninjacore.dylib\"))\n\n");
-#else
-	fprintf(out, "core = ctypes.CDLL(os.path.join(_base_path, \"libbinaryninjacore.so.1\"))\n\n");
-#endif
+	fprintf(out, "import platform\n");
+	fprintf(out, "core = None\n");
+	fprintf(out, "_base_path = None\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");
+	fprintf(out, "elif platform.system() == \"Linux\":\n");
+	fprintf(out, "\t_base_path = os.path.join(os.path.dirname(__file__), \"..\", \"..\")\n");
+	fprintf(out, "\tcore = ctypes.CDLL(os.path.join(_base_path, \"libbinaryninjacore.so.1\"))\n\n");
+	fprintf(out, "elif platform.system() == \"Windows\":\n");
+	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");
 
 	// Create type objects
 	fprintf(out, "# Type definitions\n");
-	map<string, int64_t> enumMembers;
 	for (auto& i : types)
 	{
 		string name;
@@ -191,23 +204,20 @@ int main(int argc, char* argv[])
 		if (i.second->GetClass() == StructureTypeClass)
 		{
 			fprintf(out, "class %s(ctypes.Structure):\n", name.c_str());
-			fprintf(out, "    pass\n");
+			fprintf(out, "\tpass\n");
 		}
 		else if (i.second->GetClass() == EnumerationTypeClass)
 		{
-			fprintf(out, "%s = ctypes.c_int\n", name.c_str());
-			for (auto& j : i.second->GetEnumeration()->GetMembers())
-				fprintf(out, "%s = %" PRId64 "\n", j.name.c_str(), j.value);
-			fprintf(out, "%s_names = {\n", name.c_str());
-			for (auto& j : i.second->GetEnumeration()->GetMembers())
-				fprintf(out, "    %" PRId64 ": \"%s\",\n", j.value, j.name.c_str());
-			fprintf(out, "}\n");
-			fprintf(out, "%s_by_name = {\n", name.c_str());
-			for (auto& j : i.second->GetEnumeration()->GetMembers())
-				fprintf(out, "    \"%s\": %" PRId64 ",\n", j.name.c_str(), j.value);
-			fprintf(out, "}\n");
+			if (name.size() > 2 && name.substr(0, 2) == "BN")
+				name = name.substr(2);
+
+			fprintf(out, "%sEnum = ctypes.c_int\n", name.c_str());
+
+			fprintf(enums, "\n\nclass %s(enum.IntEnum):\n", name.c_str());
 			for (auto& j : i.second->GetEnumeration()->GetMembers())
-				enumMembers[j.name] = j.value;
+			{
+				fprintf(enums, "\t%s = %" PRId64 "\n", j.name.c_str(), j.value);
+			}
 		}
 		else if ((i.second->GetClass() == BoolTypeClass) || (i.second->GetClass() == IntegerTypeClass) ||
 		         (i.second->GetClass() == FloatTypeClass) || (i.second->GetClass() == ArrayTypeClass))
@@ -218,10 +228,6 @@ int main(int argc, char* argv[])
 		}
 	}
 
-	fprintf(out, "all_enum_values = {\n");
-	for (auto& i : enumMembers)
-		fprintf(out, "    \"%s\": %" PRId64 ",\n", i.first.c_str(), i.second);
-	fprintf(out, "}\n");
 
 	fprintf(out, "\n# Structure definitions\n");
 	for (auto& i : types)
@@ -235,11 +241,11 @@ int main(int argc, char* argv[])
 			fprintf(out, "%s._fields_ = [\n", name.c_str());
 			for (auto& j : i.second->GetStructure()->GetMembers())
 			{
-				fprintf(out, "        (\"%s\", ", j.name.c_str());
+				fprintf(out, "\t\t(\"%s\", ", j.name.c_str());
 				OutputType(out, j.type);
 				fprintf(out, "),\n");
 			}
-			fprintf(out, "    ]\n");
+			fprintf(out, "\t]\n");
 		}
 	}
 
@@ -281,7 +287,7 @@ int main(int argc, char* argv[])
 			fprintf(out, "%s.argtypes = [\n", funcName.c_str());
 			for (auto& j : i.second->GetParameters())
 			{
-				fprintf(out, "        ");
+				fprintf(out, "\t\t");
 				if (name == "BNFreeString")
 				{
 					// BNFreeString expects a pointer to a string allocated by the core, so do not use
@@ -295,38 +301,39 @@ int main(int argc, char* argv[])
 				}
 				fprintf(out, ",\n");
 			}
-			fprintf(out, "    ]\n");
+			fprintf(out, "\t]\n");
 		}
 
 		if (stringResult)
 		{
 			// Emit wrapper to get Python string and free native memory
 			fprintf(out, "def %s(*args):\n", name.c_str());
-			fprintf(out, "    result = %s(*args)\n", funcName.c_str());
-			fprintf(out, "    string = ctypes.cast(result, ctypes.c_char_p).value\n");
-			fprintf(out, "    BNFreeString(result)\n");
-			fprintf(out, "    return string\n");
+			fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
+			fprintf(out, "\tstring = ctypes.cast(result, ctypes.c_char_p).value\n");
+			fprintf(out, "\tBNFreeString(result)\n");
+			fprintf(out, "\treturn string\n");
 		}
 		else if (pointerResult)
 		{
 			// Emit wrapper to return None on null pointer
 			fprintf(out, "def %s(*args):\n", name.c_str());
-			fprintf(out, "    result = %s(*args)\n", funcName.c_str());
-			fprintf(out, "    if not result:\n");
-			fprintf(out, "        return None\n");
-			fprintf(out, "    return result\n");
+			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");
 		}
 	}
 
 	fprintf(out, "\n# Helper functions\n");
 	fprintf(out, "def handle_of_type(value, handle_type):\n");
-	fprintf(out, "    if isinstance(value, ctypes.POINTER(handle_type)) or isinstance(value, ctypes.c_void_p):\n");
-	fprintf(out, "        return ctypes.cast(value, ctypes.POINTER(handle_type))\n");
-	fprintf(out, "    raise ValueError, 'expected pointer to %%s' %% str(handle_type)\n");
+	fprintf(out, "\tif isinstance(value, ctypes.POINTER(handle_type)) or isinstance(value, ctypes.c_void_p):\n");
+	fprintf(out, "\t\treturn ctypes.cast(value, ctypes.POINTER(handle_type))\n");
+	fprintf(out, "\traise ValueError, 'expected pointer to %%s' %% str(handle_type)\n");
 
 	fprintf(out, "\n# Set path for core plugins\n");
 	fprintf(out, "BNSetBundledPluginDirectory(os.path.join(_base_path, \"plugins\"))\n");
 
 	fclose(out);
+	fclose(enums);
 	return 0;
 }