feat: Implemented full script loader API in python

This commit is contained in:
WerWolv
2024-03-24 13:43:14 +01:00
parent 3144d79605
commit ccb78246ab
5 changed files with 250 additions and 61 deletions

View File

@@ -17,12 +17,7 @@ namespace hex::script::loader {
bool loadAll() override;
private:
struct Script {
void *module;
void *mainFunction;
};
std::vector<Script> m_scripts;
std::vector<void*> m_loadedModules;
};
}

View File

@@ -44,7 +44,8 @@ FUNCTION_DEFINITION(void, PyThreadState_DeleteCurrent, (), ())
FUNCTION_DEFINITION(PyThreadState *, PyThreadState_New, (PyInterpreterState *interp), (interp))
FUNCTION_DEFINITION(PyObject *, PyImport_AddModule, (const char *name), (name))
FUNCTION_DEFINITION(PyObject *, PyModule_New, (const char *name), (name))
FUNCTION_DEFINITION(PyObject *, PyObject_GetAttrString, (PyObject *pobj, const char *name), (pobj, name))
FUNCTION_DEFINITION(int, PyObject_HasAttrString, (PyObject *pobj, const char *name), (pobj, name))
bool initPythonLoader() {
void *pythonLibrary = nullptr;
@@ -97,6 +98,9 @@ bool initPythonLoader() {
INIT_FUNCTION(PyCallable_Check);
INIT_FUNCTION(PyObject_CallObject);
INIT_FUNCTION(PyObject_GetAttrString);
INIT_FUNCTION(PyObject_HasAttrString);
INIT_FUNCTION(_Py_Dealloc);

View File

@@ -41,59 +41,64 @@ namespace hex::script::loader {
return true;
}
std::string getCurrentTraceback() {
PyObject *ptype, *pvalue, *ptraceback;
PyErr_Fetch(&ptype, &pvalue, &ptraceback);
PyErr_NormalizeException(&ptype, &pvalue, &ptraceback);
namespace {
PyObject *pModuleName = PyUnicode_FromString("traceback");
PyObject *pModule = PyImport_Import(pModuleName);
Py_DECREF(pModuleName);
std::string getCurrentTraceback() {
PyObject *ptype, *pvalue, *ptraceback;
PyErr_Fetch(&ptype, &pvalue, &ptraceback);
PyErr_NormalizeException(&ptype, &pvalue, &ptraceback);
if (pModule != nullptr) {
PyObject *pDict = PyModule_GetDict(pModule);
PyObject *pFunc = PyDict_GetItemString(pDict, "format_exception");
PyObject *pModuleName = PyUnicode_FromString("traceback");
PyObject *pModule = PyImport_Import(pModuleName);
Py_DECREF(pModuleName);
if (pFunc && PyCallable_Check(pFunc)) {
PyObject *pArgs = PyTuple_New(3);
PyTuple_SetItem(pArgs, 0, ptype);
PyTuple_SetItem(pArgs, 1, pvalue);
PyTuple_SetItem(pArgs, 2, ptraceback);
if (pModule != nullptr) {
PyObject *pDict = PyModule_GetDict(pModule);
PyObject *pFunc = PyDict_GetItemString(pDict, "format_exception");
PyObject *pResult = PyObject_CallObject(pFunc, pArgs);
Py_DECREF(pArgs);
if (pFunc && PyCallable_Check(pFunc)) {
PyObject *pArgs = PyTuple_New(3);
PyTuple_SetItem(pArgs, 0, ptype);
PyTuple_SetItem(pArgs, 1, pvalue);
PyTuple_SetItem(pArgs, 2, ptraceback);
if (pResult != NULL) {
const char *errorMessage = PyUnicode_AsUTF8(PyUnicode_Join(PyUnicode_FromString(""), pResult));
Py_DECREF(pResult);
Py_DECREF(pModule);
return errorMessage;
PyObject *pResult = PyObject_CallObject(pFunc, pArgs);
Py_DECREF(pArgs);
if (pResult != NULL) {
const char *errorMessage = PyUnicode_AsUTF8(PyUnicode_Join(PyUnicode_FromString(""), pResult));
Py_DECREF(pResult);
Py_DECREF(pModule);
return errorMessage;
}
}
Py_DECREF(pModule);
}
Py_DECREF(pModule);
PyErr_Clear();
return "";
}
PyErr_Clear();
return "";
}
void populateModule(PyObject *pyModule, const std::string &sourceCode) {
PyModule_AddStringConstant(pyModule, "__file__", "");
void populateModule(PyObject *pyModule, const std::string &sourceCode) {
PyModule_AddStringConstant(pyModule, "__file__", "");
PyObject *localDict = PyModule_GetDict(pyModule);
PyObject *builtins = PyEval_GetBuiltins();
PyObject *localDict = PyModule_GetDict(pyModule);
PyObject *builtins = PyEval_GetBuiltins();
PyDict_SetItemString(localDict, "__builtins__", builtins);
PyDict_SetItemString(localDict, "__builtins__", builtins);
PyErr_Clear();
PyObject *pyValue = PyRun_String(sourceCode.c_str(), Py_file_input, localDict, localDict);
if (pyValue != nullptr) {
Py_DECREF(pyValue);
} else {
log::error("{}", getCurrentTraceback());
PyErr_Clear();
PyObject *pyValue = PyRun_String(sourceCode.c_str(), Py_file_input, localDict, localDict);
if (pyValue != nullptr) {
Py_DECREF(pyValue);
} else {
log::error("{}", getCurrentTraceback());
}
}
}
bool PythonLoader::loadAll() {
this->clearScripts();
@@ -118,26 +123,39 @@ namespace hex::script::loader {
if (!scriptFile.isValid())
continue;
this->addScript(entry.path().stem().string(), [pathString, scriptContent = scriptFile.readString()] {
PyThreadState* ts = PyThreadState_New(mainThreadState);
PyEval_RestoreThread(ts);
PyThreadState* ts = PyThreadState_New(mainThreadState);
PyEval_RestoreThread(ts);
ON_SCOPE_EXIT {
PyThreadState_Clear(ts);
PyThreadState_DeleteCurrent();
};
PyObject *libraryModule = PyImport_AddModule("imhex");
ON_SCOPE_EXIT {
PyThreadState_Clear(ts);
PyThreadState_DeleteCurrent();
};
PyModule_AddStringConstant(libraryModule, "__script_loader__", hex::format("{}", (u64)hex::getContainingModule((void*)&getCurrentTraceback)).c_str());
populateModule(libraryModule, romfs::get("python/imhex.py").data<const char>());
PyObject *libraryModule = PyImport_AddModule("imhex");
PyModule_AddStringConstant(libraryModule, "__script_loader__", hex::format("{}", reinterpret_cast<intptr_t>(hex::getContainingModule((void*)&getCurrentTraceback))).c_str());
populateModule(libraryModule, romfs::get("python/imhex.py").data<const char>());
PyObject *mainModule = PyModule_New(pathString.c_str());
populateModule(mainModule, scriptContent);
PyObject *mainModule = PyModule_New(pathString.c_str());
populateModule(mainModule, scriptFile.readString());
Py_DECREF(mainModule);
});
if (PyObject_HasAttrString(mainModule, "main")) {
this->addScript(entry.path().stem().string(), [mainModule] {
PyThreadState* ts = PyThreadState_New(mainThreadState);
PyEval_RestoreThread(ts);
ON_SCOPE_EXIT {
PyThreadState_Clear(ts);
PyThreadState_DeleteCurrent();
};
auto mainFunction = PyObject_GetAttrString(mainModule, "main");
PyObject_CallObject(mainFunction, nullptr);
Py_DECREF(mainFunction);
});
}
m_loadedModules.push_back(mainModule);
}
}

View File

@@ -26,6 +26,14 @@ SCRIPT_API(void writeMemory, u64 address, size_t size, const void *buffer) {
provider->write(address, buffer, size);
}
SCRIPT_API(u64 getBaseAddress) {
return hex::ImHexApi::Provider::get()->getBaseAddress();
}
SCRIPT_API(u64 getDataSize) {
return hex::ImHexApi::Provider::get()->getSize();
}
SCRIPT_API(bool getSelection, u64 *start, u64 *end) {
if (start == nullptr || end == nullptr)
return false;

View File

@@ -1,9 +1,173 @@
import ctypes
script_loader = ctypes.CDLL("Script Loader", ctypes.DEFAULT_MODE, int(__script_loader__))
from enum import Enum
from abc import ABC, abstractmethod
_script_loader = ctypes.CDLL("Script Loader", ctypes.DEFAULT_MODE, int(__script_loader__))
_callback_refs = []
class Color:
def __init__(self, r: int, g: int, b: int, a: int):
self.r = r
self.g = g
self.b = b
self.a = a
def to_int(self):
return (self.a << 24) | (self.b << 16) | (self.g << 8) | self.r
class UI:
@staticmethod
def show_message_box(message: str):
script_loader.showMessageBoxV1(ctypes.create_string_buffer(message.encode("utf-8")))
_script_loader.showMessageBoxV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def show_input_text_box(title: str, message: str, buffer_size: int = 256):
buffer = ctypes.create_string_buffer(buffer_size)
_script_loader.showInputTextBoxV1(ctypes.create_string_buffer(title.encode("utf-8")),
ctypes.create_string_buffer(message.encode("utf-8")),
buffer, buffer_size)
return buffer.value.decode("utf-8")
@staticmethod
def show_yes_no_question_box(title: str, message: str):
result = ctypes.c_bool()
_script_loader.showYesNoQuestionBoxV1(ctypes.create_string_buffer(title.encode("utf-8")),
ctypes.create_string_buffer(message.encode("utf-8")),
ctypes.byref(result))
class ToastType(Enum):
INFO = 0
WARNING = 1
ERROR = 2
@staticmethod
def show_toast(message: str, toast_type: ToastType):
_script_loader.showToastV1(ctypes.create_string_buffer(message.encode("utf-8")), toast_type.value)
@staticmethod
def get_imgui_context():
return _script_loader.getImGuiContextV1()
@staticmethod
def register_view(icon: str, name: str, draw_callback):
draw_function = ctypes.CFUNCTYPE(None)
global _callback_refs
_callback_refs.append(draw_function(draw_callback))
_script_loader.registerViewV1(ctypes.create_string_buffer(icon.encode("utf-8")),
ctypes.create_string_buffer(name.encode("utf-8")),
draw_function(draw_callback))
@staticmethod
def add_menu_item(icon: str, menu_name: str, item_name: str, callback):
callback_function = ctypes.CFUNCTYPE(None)
global _callback_refs
_callback_refs.append(callback_function(callback))
_script_loader.addMenuItemV1(ctypes.create_string_buffer(icon.encode("utf-8")),
ctypes.create_string_buffer(menu_name.encode("utf-8")),
ctypes.create_string_buffer(item_name.encode("utf-8")),
_callback_refs[-1])
class Bookmarks:
@staticmethod
def create_bookmark(address: int, size: int, color: Color, name: str, description: str = ""):
_script_loader.addBookmarkV1(address, size,
color.to_int(),
ctypes.create_string_buffer(name.encode("utf-8")),
ctypes.create_string_buffer(description.encode("utf-8")))
class Logger:
@staticmethod
def print(message: str):
_script_loader.logPrintV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def println(message: str):
_script_loader.logPrintlnV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def debug(message: str):
_script_loader.logDebugV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def info(message: str):
_script_loader.logInfoV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def warn(message: str):
_script_loader.logWarnV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def error(message: str):
_script_loader.logErrorV1(ctypes.create_string_buffer(message.encode("utf-8")))
@staticmethod
def fatal(message: str):
_script_loader.logFatalV1(ctypes.create_string_buffer(message.encode("utf-8")))
class Memory:
@staticmethod
def read(address: int, size: int):
buffer = ctypes.create_string_buffer(size)
_script_loader.readMemoryV1(address, buffer, size)
return buffer.raw
@staticmethod
def write(address: int, data: bytes):
_script_loader.writeMemoryV1(address, data, len(data))
@staticmethod
def get_base_address():
return _script_loader.getBaseAddressV1()
@staticmethod
def get_data_size():
return _script_loader.getDataSizeV1()
@staticmethod
def get_selection():
start = ctypes.c_uint64()
end = ctypes.c_uint64()
if not _script_loader.getSelectionV1(ctypes.byref(start), ctypes.byref(end)):
return None, None
else:
return start.value, end.value
class Provider(ABC):
def __init__(self, type_name, name):
self.type_name = type_name
self.name = name
@abstractmethod
def read(self, address: int, size: int):
pass
@abstractmethod
def write(self, address: int, data: bytes):
pass
@abstractmethod
def get_size(self):
pass
@staticmethod
def register_provider(provider):
provider_read_function = ctypes.CFUNCTYPE(None, ctypes.c_uint64, ctypes.c_void_p, ctypes.c_uint64)
provider_write_function = ctypes.CFUNCTYPE(None, ctypes.c_uint64, ctypes.c_void_p, ctypes.c_uint64)
provider_get_size_function = ctypes.CFUNCTYPE(ctypes.c_uint64)
global _callback_refs
_callback_refs.append(provider_read_function(provider.read))
_callback_refs.append(provider_write_function(provider.write))
_callback_refs.append(provider_get_size_function(provider.get_size))
_script_loader.registerMemoryProviderV1(ctypes.create_string_buffer(provider.type_name.encode("utf-8")),
ctypes.create_string_buffer(provider.name.encode("utf-8")),
_callback_refs[-3],
_callback_refs[-2],
_callback_refs[-1])