1 #include "helpers/loader_script_handler.hpp"
2 
3 #include <hex/views/view.hpp>
4 #include <hex/helpers/utils.hpp>
5 #include <hex/providers/provider.hpp>
6 
7 #define PY_SSIZE_T_CLEAN
8 #include <Python.h>
9 #include <structmember.h>
10 
11 #include <cstring>
12 #include <filesystem>
13 
14 using namespace std::literals::string_literals;
15 
16 namespace hex {
17 
Py_getFilePath(PyObject * self,PyObject * args)18     PyObject* LoaderScript::Py_getFilePath(PyObject *self, PyObject *args) {
19         return PyUnicode_FromString(LoaderScript::s_filePath.c_str());
20     }
21 
Py_addPatch(PyObject * self,PyObject * args)22     PyObject* LoaderScript::Py_addPatch(PyObject *self, PyObject *args) {
23         u64 address;
24         u8 *patches;
25         Py_ssize_t count;
26 
27         if (!PyArg_ParseTuple(args, "K|y#", &address, &patches, &count)) {
28             PyErr_BadArgument();
29             return nullptr;
30         }
31 
32         if (patches == nullptr || count == 0) {
33             PyErr_SetString(PyExc_TypeError, "Invalid patch provided");
34             return nullptr;
35         }
36 
37         if (address >= LoaderScript::s_dataProvider->getActualSize()) {
38             PyErr_SetString(PyExc_IndexError, "address out of range");
39             return nullptr;
40         }
41 
42         LoaderScript::s_dataProvider->write(address, patches, count);
43 
44         Py_RETURN_NONE;
45     }
46 
Py_addBookmark(PyObject * self,PyObject * args)47     PyObject* LoaderScript::Py_addBookmark(PyObject *self, PyObject *args) {
48         u64 address;
49         size_t size;
50 
51         char *name = nullptr;
52         char *comment = nullptr;
53 
54         if (!PyArg_ParseTuple(args, "K|n|s|s", &address, &size, &name, &comment)) {
55             PyErr_BadArgument();
56             return nullptr;
57         }
58 
59         if (name == nullptr || comment == nullptr) {
60             PyErr_SetString(PyExc_IndexError, "address out of range");
61             return nullptr;
62         }
63 
64         ImHexApi::Bookmarks::add(address, size, name, comment);
65 
66         Py_RETURN_NONE;
67     }
68 
createStructureType(std::string keyword,PyObject * args)69     static PyObject* createStructureType(std::string keyword, PyObject *args) {
70         auto type = PyTuple_GetItem(args, 0);
71         if (type == nullptr) {
72             PyErr_BadArgument();
73             return nullptr;
74         }
75 
76         auto instance = PyObject_CallObject(type, nullptr);
77         if (instance == nullptr) {
78             PyErr_BadArgument();
79             return nullptr;
80         }
81 
82         SCOPE_EXIT( Py_DECREF(instance); );
83 
84         if (instance->ob_type->tp_base == nullptr || instance->ob_type->tp_base->tp_name != "ImHexType"s) {
85             PyErr_SetString(PyExc_TypeError, "class type must extend from ImHexType");
86             return nullptr;
87         }
88 
89         auto dict = instance->ob_type->tp_dict;
90         if (dict == nullptr) {
91             PyErr_BadArgument();
92             return nullptr;
93         }
94 
95         auto annotations = PyDict_GetItemString(dict, "__annotations__");
96         if (annotations == nullptr) {
97             PyErr_BadArgument();
98             return nullptr;
99         }
100 
101         auto list = PyDict_Items(annotations);
102         if (list == nullptr) {
103             PyErr_BadArgument();
104             return nullptr;
105         }
106 
107         SCOPE_EXIT( Py_DECREF(list); );
108 
109         std::string code = keyword + " " + instance->ob_type->tp_name + " {\n";
110 
111         for (u16 i = 0; i < PyList_Size(list); i++) {
112             auto item = PyList_GetItem(list, i);
113 
114             auto memberName = PyUnicode_AsUTF8(PyTuple_GetItem(item, 0));
115             auto memberType = PyTuple_GetItem(item, 1);
116             if (memberType == nullptr) {
117                 PyErr_SetString(PyExc_TypeError, "member needs to have a annotation extending from ImHexType");
118                 return nullptr;
119             }
120 
121             // Array already is an object
122             if (memberType->ob_type->tp_name == "array"s) {
123 
124                 auto arrayType = PyObject_GetAttrString(memberType, "array_type");
125                 if (arrayType == nullptr) {
126                     PyErr_BadArgument();
127                     return nullptr;
128                 }
129 
130                 code += "   "s + arrayType->ob_type->tp_name + " " + memberName;
131 
132                 auto arraySize = PyObject_GetAttrString(memberType, "size");
133                 if (arraySize == nullptr) {
134                     PyErr_BadArgument();
135                     return nullptr;
136                 }
137 
138                 if (PyUnicode_Check(arraySize))
139                     code += "["s + PyUnicode_AsUTF8(arraySize) + "];\n";
140                 else if (PyLong_Check(arraySize))
141                     code += "["s + std::to_string(PyLong_AsLong(arraySize)) + "];\n";
142                 else {
143                     PyErr_SetString(PyExc_TypeError, "invalid array size type. Expected string or int");
144                     return nullptr;
145                 }
146 
147 
148             } else {
149                 auto memberTypeInstance = PyObject_CallObject(memberType, nullptr);
150                 if (memberTypeInstance == nullptr || memberTypeInstance->ob_type->tp_base == nullptr || memberTypeInstance->ob_type->tp_base->tp_name != "ImHexType"s) {
151                     PyErr_SetString(PyExc_TypeError, "member needs to have a annotation extending from ImHexType");
152                     if (memberTypeInstance != nullptr)
153                         Py_DECREF(memberTypeInstance);
154 
155                     return nullptr;
156                 }
157 
158                 code += "   "s + memberTypeInstance->ob_type->tp_name + " "s + memberName + ";\n";
159 
160                 Py_DECREF(memberTypeInstance);
161             }
162         }
163 
164         code += "};\n";
165 
166         View::postEvent(Events::AppendPatternLanguageCode, code.c_str());
167 
168         Py_RETURN_NONE;
169     }
170 
Py_addStruct(PyObject * self,PyObject * args)171     PyObject* LoaderScript::Py_addStruct(PyObject *self, PyObject *args) {
172         return createStructureType("struct", args);
173     }
174 
Py_addUnion(PyObject * self,PyObject * args)175     PyObject* LoaderScript::Py_addUnion(PyObject *self, PyObject *args) {
176         return createStructureType("union", args);
177     }
178 
processFile(std::string_view scriptPath)179     bool LoaderScript::processFile(std::string_view scriptPath) {
180         Py_SetProgramName(Py_DecodeLocale((SharedData::mainArgv)[0], nullptr));
181 
182         if (std::filesystem::exists(std::filesystem::path((SharedData::mainArgv)[0]).parent_path().string() + "/lib/python" PYTHON_VERSION_MAJOR_MINOR))
183             Py_SetPythonHome(Py_DecodeLocale(std::filesystem::path((SharedData::mainArgv)[0]).parent_path().string().c_str(), nullptr));
184 
185         PyImport_AppendInittab("_imhex", []() -> PyObject* {
186 
187             static PyMethodDef ImHexMethods[] = {
188                 { "get_file_path",  &LoaderScript::Py_getFilePath,  METH_NOARGS,  "Returns the path of the file being loaded."  },
189                 { "patch",          &LoaderScript::Py_addPatch,     METH_VARARGS, "Patches a region of memory"                  },
190                 { "add_bookmark",   &LoaderScript::Py_addBookmark,  METH_VARARGS, "Adds a bookmark"                             },
191                 { "add_struct",     &LoaderScript::Py_addStruct,    METH_VARARGS, "Adds a struct"                               },
192                 { "add_union",      &LoaderScript::Py_addUnion,     METH_VARARGS, "Adds a union"                                },
193                 { nullptr,          nullptr,               0,     nullptr                                       }
194             };
195 
196             static PyModuleDef ImHexModule = {
197                 PyModuleDef_HEAD_INIT, "imhex", nullptr, -1, ImHexMethods, nullptr, nullptr, nullptr, nullptr
198             };
199 
200             auto module =  PyModule_Create(&ImHexModule);
201             if (module == nullptr)
202                 return nullptr;
203 
204             return module;
205         });
206 
207         Py_Initialize();
208 
209         {
210             auto sysPath = PySys_GetObject("path");
211             auto path = PyUnicode_FromString("lib");
212 
213             PyList_Insert(sysPath, 0, path);
214         }
215 
216         FILE *scriptFile = fopen(scriptPath.data(), "r");
217         PyRun_SimpleFile(scriptFile, scriptPath.data());
218 
219         fclose(scriptFile);
220 
221         Py_Finalize();
222 
223         return true;
224     }
225 
226 }