Support cs_disasm_iter in Python API binding about capstone HOT 8 CLOSED

#include <Python.h>
#include <capstone/capstone.h>

// A struct to hold the Capstone handle and the current instruction
typedef struct {
    csh handle;
    cs_insn *insn;
} DisassemblerObject;

// The __del__ method of the Disassembler class
static void Disassembler_dealloc(DisassemblerObject* self)
{
    cs_free(self->insn, 1);
    cs_close(&self->handle);
    Py_TYPE(self)->tp_free((PyObject*)self);
}

// The __init__ method of the Disassembler class
static int Disassembler_init(DisassemblerObject* self, PyObject* args)
{
    // Initialize the Capstone handle
    if (cs_open(CS_ARCH_X86, CS_MODE_64, &self->handle) != CS_ERR_OK) {
        return -1;
    }

    // Allocate memory for the current instruction
    self->insn = cs_malloc(self->handle);
    if (!self->insn) {
        cs_close(&self->handle);
        return -1;
    }

    return 0;
}

// The disasm method of the Disassembler class
static PyObject* Disassembler_disasm(DisassemblerObject* self, PyObject* args)
{
    const uint8_t* code;
    size_t size;
    uint64_t address;

    // Parse the Python arguments
    if (!PyArg_ParseTuple(args, "y#K", &code, &size, &address)) {
        return NULL;
    }

    // Use cs_disasm_iter to disassemble the next instruction
    if (!cs_disasm_iter(self->handle, &code, &size, &address, self->insn)) {
        return NULL;
    }

    // Create a Python tuple to return
    PyObject* result = Py_BuildValue("(s,K)", self->insn->mnemonic, self->insn->address);
    return result;
}

// The methods of the Disassembler class
static PyMethodDef Disassembler_methods[] = {
    {"disasm", (PyCFunction)Disassembler_disasm, METH_VARARGS, "Disassemble the next instruction"},
    {NULL}  // Sentinel
};

// The Disassembler type object
static PyTypeObject DisassemblerType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "capstone.Disassembler",
    .tp_doc = "Disassembler objects",
    .tp_basicsize = sizeof(DisassemblerObject),
    .tp_itemsize = 0,
    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .tp_new = PyType_GenericNew,
    .tp_init = (initproc) Disassembler_init,
    .tp_dealloc = (destructor) Disassembler_dealloc,
    .tp_methods = Disassembler_methods,
};

// The module's initialization function
PyMODINIT_FUNC PyInit_capstone(void)
{
    PyObject* m;

    if (PyType_Ready(&DisassemblerType) < 0) {
        return NULL;
    }

    m = PyModule_Create(&moduledef);
    if (m == NULL) {
        return NULL;
    }

    Py_INCREF(&DisassemblerType);
    PyModule_AddObject(m, "Disassembler", (PyObject *)&DisassemblerType);

    return m;
}

diff --git a/sqlelf/elf/instruction.py b/sqlelf/elf/instruction.py
index 68acf50..9a90706 100644
--- a/sqlelf/elf/instruction.py
+++ b/sqlelf/elf/instruction.py
@@ -8,6 +8,18 @@ import apsw.ext
 import capstone
 import lief
 
+import ctypes
+
+# Define the cs_disasm_iter function
+cs_disasm_iter = capstone._cs.cs_disasm_iter
+cs_disasm_iter.argtypes = [
+    ctypes.POINTER(ctypes.c_size_t),
+    ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)),
+    ctypes.POINTER(ctypes.c_size_t),
+    ctypes.POINTER(ctypes.c_uint64),
+    ctypes.POINTER(capstone._cs_insn),
+]
+
 
 def elf_instructions(binaries: list[lief.Binary]):
     def generator() -> Iterator[dict[str, Any]]:
@@ -28,6 +40,23 @@ def elf_instructions(binaries: list[lief.Binary]):
                     # super important that these accessors are pulled out
                     # of the tight loop as they can be costly
                     section_name = section.name
+
+                    # Allocate a cs_insn struct
+                    code_array = (ctypes.c_ubyte * len(data))(*data)
+                    code_array = ctypes.cast(code_array, ctypes.POINTER(ctypes.c_ubyte))
+                    insn = capstone._cs_insn()
+                    size = ctypes.c_size_t(len(data))
+                    address = ctypes.c_uint64(section.virtual_address)
+                    while cs_disasm_iter(
+                        md.csh,
+                        ctypes.byref(code_array),
+                        ctypes.byref(size),
+                        ctypes.byref(address),
+                        ctypes.byref(insn),
+                    ):
+                        print("here")
+                        print(insn)
+
                     for (address, size, mnemonic, op_str) in md.disasm_lite(
                         data, section.virtual_address
                     ):

diff --git a/sqlelf/elf/instruction.py b/sqlelf/elf/instruction.py
index 68acf50..5325b27 100644
--- a/sqlelf/elf/instruction.py
+++ b/sqlelf/elf/instruction.py
@@ -8,6 +8,18 @@ import apsw.ext
 import capstone
 import lief
 
+import ctypes
+
+# Define the cs_disasm_iter function
+cs_disasm_iter = capstone._cs.cs_disasm_iter
+cs_disasm_iter.argtypes = [
+    ctypes.c_size_t,
+    ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)),
+    ctypes.POINTER(ctypes.c_size_t),
+    ctypes.POINTER(ctypes.c_uint64),
+    ctypes.POINTER(capstone._cs_insn),
+]
+
 
 def elf_instructions(binaries: list[lief.Binary]):
     def generator() -> Iterator[dict[str, Any]]:
@@ -28,15 +40,26 @@ def elf_instructions(binaries: list[lief.Binary]):
                     # super important that these accessors are pulled out
                     # of the tight loop as they can be costly
                     section_name = section.name
-                    for (address, size, mnemonic, op_str) in md.disasm_lite(
-                        data, section.virtual_address
+
+                    # Allocate a cs_insn struct
+                    code_array = (ctypes.c_ubyte * len(data))()
+                    code_array = ctypes.cast(code_array, ctypes.POINTER(ctypes.c_ubyte))
+                    insn = capstone._cs_insn()
+                    size = ctypes.c_size_t(len(data))
+                    address = ctypes.c_uint64(section.virtual_address)
+                    while cs_disasm_iter(
+                        md.csh,
+                        ctypes.byref(code_array),
+                        ctypes.byref(size),
+                        ctypes.byref(address),
+                        ctypes.byref(insn),
                     ):
                         yield {
                             "path": binary_name,
                             "section": section_name,
-                            "mnemonic": mnemonic,
+                            "mnemonic": insn.mnemonic.decode("ascii"),
                             "address": address,
-                            "operands": op_str,
+                            "operands": insn.op_str.decode("ascii"),
                         }
 
     return generator