From ccb78246abbf8bfad2e30c22f8275af3c73d0331 Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Sun, 24 Mar 2024 13:43:14 +0100
Subject: [PATCH] feat: Implemented full script loader API in python

---
 .../include/loaders/python/python_loader.hpp  |   7 +-
 .../source/loaders/python/library_wrapper.cpp |   6 +-
 .../source/loaders/python/python_loader.cpp   | 122 +++++++------
 .../source/script_api/v1/mem.cpp              |   8 +
 plugins/script_loader/support/python/imhex.py | 168 +++++++++++++++++-
 5 files changed, 250 insertions(+), 61 deletions(-)

diff --git a/plugins/script_loader/include/loaders/python/python_loader.hpp b/plugins/script_loader/include/loaders/python/python_loader.hpp
index 1fc2fbc5e..e84133bef 100644
--- a/plugins/script_loader/include/loaders/python/python_loader.hpp
+++ b/plugins/script_loader/include/loaders/python/python_loader.hpp
@@ -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;
     };
 
 }
\ No newline at end of file
diff --git a/plugins/script_loader/source/loaders/python/library_wrapper.cpp b/plugins/script_loader/source/loaders/python/library_wrapper.cpp
index 415b01170..766f5a683 100644
--- a/plugins/script_loader/source/loaders/python/library_wrapper.cpp
+++ b/plugins/script_loader/source/loaders/python/library_wrapper.cpp
@@ -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);
 
diff --git a/plugins/script_loader/source/loaders/python/python_loader.cpp b/plugins/script_loader/source/loaders/python/python_loader.cpp
index 0370e2b21..920efd3da 100644
--- a/plugins/script_loader/source/loaders/python/python_loader.cpp
+++ b/plugins/script_loader/source/loaders/python/python_loader.cpp
@@ -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);
             }
         }
 
diff --git a/plugins/script_loader/source/script_api/v1/mem.cpp b/plugins/script_loader/source/script_api/v1/mem.cpp
index 16fcd7a78..24ba00e8a 100644
--- a/plugins/script_loader/source/script_api/v1/mem.cpp
+++ b/plugins/script_loader/source/script_api/v1/mem.cpp
@@ -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;
diff --git a/plugins/script_loader/support/python/imhex.py b/plugins/script_loader/support/python/imhex.py
index d91ee7a7f..5ee808525 100644
--- a/plugins/script_loader/support/python/imhex.py
+++ b/plugins/script_loader/support/python/imhex.py
@@ -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])
\ No newline at end of file