csonpath_python.c

#include <Python.h>
#include <stdio.h>
#include <stdlib.h>

#define CSONPATH_MAX_ARRAY_PAD (1 << 16)

#define CSONPATH_JSON PyObject *

#define CSONPATH_NULL Py_None

#define CSONPATH_GET csonpath_python_get

#define CSONPATH_AT csonpath_python_at

#define CSONPATH_REMOVE(o) Py_XDECREF(o)

#define CSONPATH_INREF(o) Py_INCREF(o)

#define CSONPATH_NEW_OBJECT() PyDict_New()

#define CSONPATH_NEW_ARRAY() PyList_New(0)

#define CSONPATH_NEW_STR(s) PyUnicode_FromString(s)

#define CSONPATH_NEW_INT(n) PyLong_FromLong(n)

#define CSONPATH_NEW_BOOL(b) PyBool_FromLong((b) ? 1 : 0)

#define CSONPATH_IS_OBJ(o) PyDict_Check(o)

#define CSONPATH_IS_ARRAY(o) PyList_Check(o)

#define CSONPATH_IS_STR(o) PyUnicode_Check(o)

#define CSONPATH_IS_NUM(o) PyLong_Check(o)

#define CSONPATH_IS_NULL(o) (o == Py_None)

#define CSONPATH_IS_BOOL(o) PyBool_Check(o)

/* that one assume CSONPATH_IS_BOOL have been called before calling CSONPATH_GET_BOOL */
#define CSONPATH_GET_BOOL(o) (o == Py_True)

#define CSONPATH_GET_STR(obj)			\
    PyUnicode_AsUTF8(obj)

#define CSONPATH_GET_NUM(obj)			\
    PyLong_AsLong(obj)

#define CSONPATH_EQUAL_NUM(obj, to_cmp)		\
    ({						\
	_Bool r = 0;				\
	if (PyLong_Check(obj)) {		\
	    r = PyLong_AsLong(obj) == to_cmp;	\
	}					\
	r;					\
    })

#define CSONPATH_EQUAL_STR(obj, to_cmp)	({			\
    _Bool r = 0;						\
      if (PyUnicode_Check(obj))	{				\
	const char *py_str = PyUnicode_AsUTF8(obj);		\
	r = !strcmp(py_str, to_cmp);				\
      }								\
      r;							\
    })

#define CSONPATH_FORMAT_EXCEPTION(ARGS...)	\
    PyErr_Format(PyExc_ValueError, ARGS)

#define CSONPATH_EXCEPTION(ARGS...)		\
    PyErr_Format(PyExc_ValueError, ARGS);	\
    return -1;

#define CSONPATH_CALLBACK PyObject *

#define CSONPATH_CALLBACK_DATA PyObject *

#define CSONPATH_CALL_CALLBACK(callback, ctx, child_info, tmp, udata) \
    ({								\
	int _cb_ok = 0;							\
	PyObject *arglist;						\
	PyObject *_cb_ret;						\
	if (child_info->type == CSONPATH_STR)				\
	    arglist = Py_BuildValue("(OsOO)", ctx, child_info->key, tmp, udata); \
	else								\
	    arglist = Py_BuildValue("(OiOO)", ctx, child_info->idx, tmp, udata); \
	_cb_ret = PyObject_CallObject(callback, arglist);		\
	Py_DECREF(arglist);						\
	if (!_cb_ret)							\
	    _cb_ok = -1;						\
	else								\
	    Py_DECREF(_cb_ret);						\
	_cb_ok;								\
    })

/* assuming each modification of the object need to go out of the loop */
#define CSONPATH_NEED_FOREACH_REDO(o)	1

#define CSONPATH_REMOVE_CHILD(obj, child_info)				\
  if (child_info.type == CSONPATH_INTEGER) {				\
    PyList_SetSlice(obj, child_info.idx, child_info.idx + 1, NULL);	\
  } else if (child_info.type == CSONPATH_STR) {				\
    PyDict_DelItemString(obj, child_info.key);				\
  }

#define CSONPATH_ARRAY_APPEND(list, el) PyList_Append(list, el)


static int pydict_try_setitemstring(PyObject *obj, const char * const at, PyObject *el, int do_incref)
{
    (void)do_incref;
    if (!PyDict_Check(obj)) {
	PyObject* repr = PyObject_Repr(obj);
	PyObject* str = PyUnicode_AsEncodedString(repr, "utf-8", "~E~");
	const char *bytes = PyBytes_AS_STRING(str);

	PyErr_Format(PyExc_ValueError, "Unable to follow path (%s) in %s: Dict expected", at, bytes);
	Py_XDECREF(repr);
	Py_XDECREF(str);
	return -1;
    }
    if (!at) {
        PyErr_SetString(PyExc_TypeError, "dict keys must be strings");
        return -1;
    }
    if (PyDict_SetItemString(obj, at, el) < 0)
	return -1;
    return 1;
}

static int python_set_or_insert_item(PyObject *array, Py_ssize_t at, PyObject *el, int do_incref)
{
    (void)do_incref;
    if (!PyList_Check(array)) {
	PyErr_Format(PyExc_ValueError, "Unable to follow path: List expected");
	return -1;
    }
    Py_ssize_t s = PyList_Size(array);
    if (at >= s) {
	if (at - s > CSONPATH_MAX_ARRAY_PAD) {
	    PyErr_SetString(PyExc_IndexError,
			    "array index out of range (None padding too large)");
	    return -1;
	}
	for (;s < at; ++s)
	    if (PyList_Insert(array, s, Py_None) < 0)
		return -1;

	if (PyList_Insert(array, at, el) < 0)
	    return -1;
    } else {
	Py_INCREF(el);
	if (PyList_SetItem(array, at, el) < 0) {
	    Py_DECREF(el);
	    return -1;
	}
    }
    return 1;
}

#define CSONPATH_APPEND_AT(array, at, el, do_incref)			\
    _Generic((at),							\
	     int: python_set_or_insert_item,				\
	     unsigned int: python_set_or_insert_item,			\
	     long: python_set_or_insert_item,				\
	     unsigned long: python_set_or_insert_item,			\
	     long long: python_set_or_insert_item,			\
	     unsigned long long: python_set_or_insert_item,		\
	     const char *: pydict_try_setitemstring,			\
	     char *: pydict_try_setitemstring				\
	) (array, at, el, do_incref)


#ifdef __TINYC__
#define CSONPATH_PRAGMA(...)
#else
#define CSONPATH_PRAGMA(...) _Pragma(__VA_ARGS__)
#endif


#if PY_MAJOR_VERSION > 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 13)
/* PyList_Clear was introduced in Python 3.13 */
#define CSONPATH_ARRAY_CLEAR(o) PyList_Clear(o)
#else
/* For older versions, manually clear the list */
#define CSONPATH_ARRAY_CLEAR(o)            \
    do {                                   \
        while (PyList_Size(o) > 0) {       \
            PyList_SetSlice(o, 0, PyList_Size(o), NULL); \
        }                                  \
    } while (0)
#endif

#define CSONPATH_OBJ_CLEAR(o) PyDict_Clear(o)

#define CSONPATH_ARRAY_LENGTH(o) ((size_t)PyList_Size(o))

#define CSONPATH_FOREACH_ARRAY(obj, child, idx)				\
    CSONPATH_PRAGMA("GCC unroll 8")					\
    for (intptr_t array_len = PyList_Size(obj), idx = 0;		\
	 ({int r = idx < array_len; if (r) child = PyList_GetItem(obj, idx); r;}); ++idx)

#define CSONPATH_FOREACH_OBJ(obj, child, key)				\
    for (Py_ssize_t pos_ = 0; ({					\
		PyObject *key_;						\
		_Bool r = PyDict_Next(obj, &pos_, &key_, &child);	\
		if (r) {						\
		    key = PyUnicode_AsUTF8AndSize(key_, NULL);		\
		    if (!key)						\
			r = 0;						\
		}							\
		r;							\
	    });)

#define CSONPATH_FOREACH_EXT(obj, el, code, key_idx)			\
  if (PyDict_Check(obj)) {						\
    PyObject *key_;							\
    Py_ssize_t pos_ = 0;						\
    while (PyDict_Next(obj, &pos_, &key_, &el)) {			\
      const char *key_idx = PyUnicode_AsUTF8AndSize(key_, NULL);	\
      if (!key_idx) {							\
        break;								\
      }									\
      (void)key_idx;							\
      code								\
	}								\
  } else if (PyList_Check(obj)) {					\
    int array_len_ = PyList_Size(obj);					\
    CSONPATH_PRAGMA("GCC unroll 8")					\
    for (intptr_t key_idx = 0; key_idx < array_len_; ++key_idx) {	\
      el = PyList_GetItem(obj, key_idx);				\
      if (!el) {							\
        PyErr_SetString(PyExc_RuntimeError,				\
                        "list was modified during iteration");		\
        break;								\
      }									\
      code								\
	}								\
  }


static PyObject *csonpath_python_get(PyObject *obj, const char *key) {
    if (PyDict_Check(obj)) {
        PyObject *value = PyDict_GetItemString(obj, key);
        return value ? value : Py_None;
    }
    return Py_None;
}

static PyObject *csonpath_python_at(PyObject *obj, int at) {
    if (PyList_Check(obj)) {
        if (at >= 0 && at < PyList_Size(obj)) {
            PyObject *item = PyList_GetItem(obj, at);
            return item ? item : Py_None;
        }
    }
    return Py_None;
}

#include "csonpath.h"
#include "csonpath_my_fuzzing.h"

typedef struct {
    PyObject_HEAD
    struct csonpath *cp;
} PyCsonPathObject;

typedef struct {
    PyObject_HEAD
    struct csonpath_fuzzer *fuzzer;
} PyCsonPathFuzzerObject;

static int csonpath_python_set_extra_roots(struct csonpath *cp, PyObject *extra_objs)
{
    if (extra_objs == NULL || extra_objs == Py_None)
	return 0;
    if (!PyList_Check(extra_objs) && !PyTuple_Check(extra_objs)) {
	PyErr_SetString(PyExc_TypeError, "extra_objs must be a list or tuple");
	return -1;
    }
    PyObject *list = PySequence_List(extra_objs);
    if (!list)
	return -1;
    cp->extra_roots = list;
    return 0;
}

static void csonpath_python_clear_extra_roots(struct csonpath *cp)
{
    if (cp->extra_roots != CSONPATH_NULL) {
	Py_DECREF(cp->extra_roots);
	cp->extra_roots = CSONPATH_NULL;
    }
}

#define BAD_ARG() ({fprintf(stderr, "bad argument\n"); PyErr_BadArgument(); return NULL;})


static PyObject *PyCsonPath_new(PyTypeObject *subtype, PyObject* args,
				PyObject* kwargs)
{
	static char *kwlist[] = {"path", "return_empty_array", "jq_like", NULL};
	PyCsonPathObject *self = (PyCsonPathObject *)subtype->tp_alloc(subtype, 0);
	const char *s;
	int return_empty_array = 0;
	int jq_like = 0;
	struct csonpath *ret = NULL;
	PyObject *py_ret = NULL;

	if (!self)
		BAD_ARG();

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|pp", kwlist,
					 &s, &return_empty_array, &jq_like)) {
		goto error;
					 }

	int flags = CSONPATH_NO_DETROY;
	if (return_empty_array)
		flags |= CSONPATH_RETURN_EMPTY_ARRAY;

	ret = csonpath_new_ex(s, flags);
	if (!ret) {
		PyErr_NoMemory();
		goto error;
	}
	if (ret->compile_error) {
		PyErr_Format(PyExc_ValueError, "compilation fail %s",
			     ret->compile_error);
		csonpath_destroy(ret);
		goto error;
	}

	self->cp = ret;
	return (PyObject *)self;

	error:
		Py_DECREF(self);
	return py_ret;
}

static PyObject *find_all(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
    PyObject *json;
    PyObject *extra_objs = Py_None;
    static char *kwlist[] = {"json", "extra_objs", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O", kwlist,
				     &json, &extra_objs))
        BAD_ARG();

    if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
	return NULL;

    PyObject *ret = csonpath_find_all(self->cp, json);
    csonpath_python_clear_extra_roots(self->cp);
    if (PyErr_Occurred())
        return NULL;

    return ret;
}

static PyObject *find_first(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
    PyObject *json;
    PyObject *extra_objs = Py_None;
    static char *kwlist[] = {"json", "extra_objs", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O", kwlist,
				     &json, &extra_objs))
        BAD_ARG();

    if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
	return NULL;

    PyObject *ret = csonpath_find_first(self->cp, json);
    csonpath_python_clear_extra_roots(self->cp);
    if (PyErr_Occurred())
        return NULL;

    Py_INCREF(ret);
    return ret;
}

static PyObject *print_instructions(PyCsonPathObject *self, PyObject *args, PyObject *kwds)
{
    if (!self->cp) {
        PyErr_SetString(PyExc_RuntimeError, "compiled path is NULL");
        return NULL;
    }
    csonpath_print_instruction(self->cp);
    Py_RETURN_NONE;
}

static PyObject *callback(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
    PyObject *json, *callback, *udata = Py_None;
    PyObject *extra_objs = Py_None;
    static char *kwlist[] = {"json", "callback", "udata", "extra_objs", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|OO", kwlist,
				     &json, &callback, &udata, &extra_objs))
	BAD_ARG();
    if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
	return NULL;
    int ret = csonpath_callback(self->cp, json, callback, udata);
    csonpath_python_clear_extra_roots(self->cp);
    if (PyErr_Occurred() || ret < 0) {
        return NULL;
    }
    return PyLong_FromLong(ret);
}

static PyObject *do_remove(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
  PyObject *json;
  PyObject *extra_objs = Py_None;
  static char *kwlist[] = {"json", "extra_objs", NULL};

  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O", kwlist,
				   &json, &extra_objs))
    BAD_ARG();
  if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
      return NULL;
  int ret = csonpath_remove(self->cp, json);
  csonpath_python_clear_extra_roots(self->cp);
  if (PyErr_Occurred() || ret < 0)
      return NULL;
  return PyLong_FromLong(ret);
}

static PyObject *update_or_create(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
  PyObject *json;
  PyObject *value;
  PyObject *extra_objs = Py_None;
  static char *kwlist[] = {"json", "value", "extra_objs", NULL};

  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|$O", kwlist,
				   &json, &value, &extra_objs))
      BAD_ARG();
  if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
      return NULL;
  int ret = csonpath_update_or_create(self->cp, json, value);
  csonpath_python_clear_extra_roots(self->cp);
  if (PyErr_Occurred() || ret < 0) {
      return NULL;
  }
  return PyLong_FromLong(ret);
}

static PyObject *update_or_create_callback(PyCsonPathObject *self, PyObject* args, PyObject* kwargs)
{
    PyObject *json, *callback, *udata = Py_None;
    PyObject *extra_objs = Py_None;
    static char *kwlist[] = {"json", "callback", "udata", "extra_objs", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|OO", kwlist,
				     &json, &callback, &udata, &extra_objs))
	BAD_ARG();
    if (csonpath_python_set_extra_roots(self->cp, extra_objs) < 0)
	return NULL;
    int ret = csonpath_update_or_create_callback(self->cp, json, callback, udata);
    csonpath_python_clear_extra_roots(self->cp);
    if (PyErr_Occurred() || ret < 0) {
        return NULL;
    }
    return PyLong_FromLong(ret);
}

static void PyCsonPath_dealloc(PyCsonPathObject *self) {
  if (self->cp) {
    csonpath_destroy(self->cp);
    self->cp = NULL;
  }
  Py_TYPE(self)->tp_free((PyObject *)self);
}

static PyObject *PyCsonPathFuzzer_new(PyTypeObject *subtype, PyObject* args,
				      PyObject* kwargs)
{
	static char *kwlist[] = {"paths", "seed", "obj", "skip_objects", "dry_run", NULL};
	PyCsonPathFuzzerObject *self = (PyCsonPathFuzzerObject *)subtype->tp_alloc(subtype, 0);
	PyObject *paths;
	PyObject *obj;
	PyObject *skip_objects = NULL;
	unsigned int seed = 0;
	int dry_run = 0;
	struct csonpath_fuzzer *fuzzer = NULL;

	if (!self)
		BAD_ARG();

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OIO|Op", kwlist,
					 &paths, &seed, &obj, &skip_objects, &dry_run)) {
		goto error;
	}

	if (!PyList_Check(paths)) {
		PyErr_SetString(PyExc_TypeError, "paths must be a list");
		goto error;
	}

	fuzzer = csonpath_fuzzer_new(paths, seed, obj);
	if (!fuzzer) {
		PyErr_SetString(PyExc_ValueError, "failed to create fuzzer");
		goto error;
	}

	fuzzer->dry_run = dry_run;

	if (skip_objects && skip_objects != Py_None) {
		if (!PyList_Check(skip_objects)) {
			PyErr_SetString(PyExc_TypeError, "skip_objects must be a list");
			goto error;
		}
		Py_ssize_t n = PyList_Size(skip_objects);
		for (Py_ssize_t i = 0; i < n; ++i) {
			PyObject *target = PyList_GetItem(skip_objects, i);
			csonpath_fuzzer_mark_seen(fuzzer, target);
		}
	}

	self->fuzzer = fuzzer;
	return (PyObject *)self;

 error:
	Py_DECREF(self);
	return NULL;
}

static PyObject *PyCsonPathFuzzer_step(PyCsonPathFuzzerObject *self, PyObject* Py_UNUSED(args))
{
	enum csonpath_fuzzer_action action = csonpath_fuzz_step(self->fuzzer);
	return PyLong_FromLong((long)action);
}

static void PyCsonPathFuzzer_dealloc(PyCsonPathFuzzerObject *self) {
	if (self->fuzzer) {
		csonpath_fuzzer_destroy(self->fuzzer);
		self->fuzzer = NULL;
	}
	Py_TYPE(self)->tp_free((PyObject *)self);
}


static PyObject *PyCsonPath_set_path(PyCsonPathObject *self, PyObject* args) {
    const char *new_path;
    if (!PyArg_ParseTuple(args, "s", &new_path))
      return NULL;
    if (!new_path) return Py_False;

    struct csonpath *new_cjp = csonpath_new_ex(new_path, CSONPATH_NO_DETROY);
    if (!new_cjp) {
        PyErr_NoMemory();
        return NULL;
    }
    if (new_cjp->compile_error) {
        PyErr_Format(PyExc_ValueError, "compilation fail %s",
                     new_cjp->compile_error);
        csonpath_destroy(new_cjp);
        return NULL;
    }
    csonpath_destroy(self->cp);
    self->cp = new_cjp;
    return Py_True;
}

static PyMethodDef csonpath_fuzzer_py_method[] = {
    {"step", (PyCFunction)PyCsonPathFuzzer_step, METH_VARARGS, "apply one fuzz step"},
    {NULL, NULL, 0, NULL}
};

static PyTypeObject PyCsonPathFuzzerType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "csonpath.Fuzzer",
    .tp_basicsize = sizeof(PyCsonPathFuzzerObject),
    .tp_dealloc = (destructor)PyCsonPathFuzzer_dealloc,
    .tp_flags = Py_TPFLAGS_DEFAULT,
    .tp_methods = csonpath_fuzzer_py_method,
    .tp_new = PyCsonPathFuzzer_new
};

static PyMethodDef csonpath_py_method[] = {
    {"set_path", (PyCFunction)PyCsonPath_set_path, METH_VARARGS, "set_path"},
    {"callback", (PyCFunction)callback, METH_VARARGS | METH_KEYWORDS, "callback"},
    {"print_instructions", (PyCFunction)print_instructions, METH_NOARGS, "print_instructions"},
    {"update_or_create_callback", (PyCFunction)update_or_create_callback, METH_VARARGS | METH_KEYWORDS, "update_or_create_callback"},
    {"find_first", (PyCFunction)find_first, METH_VARARGS | METH_KEYWORDS, "find first elems"},
    {"find_all", (PyCFunction)find_all, METH_VARARGS | METH_KEYWORDS, "find all elems, if one found, pout it in an array"},
    {"remove", (PyCFunction)do_remove, METH_VARARGS | METH_KEYWORDS, "remove all elems found"},
    {"update_or_create", (PyCFunction)update_or_create, METH_VARARGS | METH_KEYWORDS, "update or create"},
    {NULL, NULL, 0, NULL}
};

static PyTypeObject PyCsonPathType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "csonpath.CsonPath",
    .tp_basicsize = sizeof(PyCsonPathObject),
    .tp_dealloc = (destructor)PyCsonPath_dealloc,
    .tp_flags = Py_TPFLAGS_DEFAULT,
    .tp_methods = csonpath_py_method,
    .tp_new = PyCsonPath_new
};

static struct PyModuleDef csonpath_py_mod = {
    PyModuleDef_HEAD_INIT,
    "csonpath",
    NULL,
    -1,
    NULL
};

PyMODINIT_FUNC PyInit_csonpath(void) {
  PyObject *m;
  if (PyType_Ready(&PyCsonPathType) < 0)
    return NULL;
  if (PyType_Ready(&PyCsonPathFuzzerType) < 0)
    return NULL;

  m = PyModule_Create(&csonpath_py_mod);
  if (!m) return NULL;

  PyModule_AddObject(m, "CsonPath", (PyObject *)&PyCsonPathType);
  PyModule_AddObject(m, "Fuzzer", (PyObject *)&PyCsonPathFuzzerType);

  PyModule_AddIntConstant(m, "MODIFY_STR", CSONPATH_FUZZER_MODIFY_STR);
  PyModule_AddIntConstant(m, "MODIFY_CNT_TYPE", CSONPATH_FUZZER_MODIFY_CNT_TYPE);
  PyModule_AddIntConstant(m, "MODIFY_STR_TYPE", CSONPATH_FUZZER_MODIFY_STR_TYPE);
  PyModule_AddIntConstant(m, "NO_ACTION_LEFT", CSONPATH_FUZZER_NO_ACTION_LEFT);

  return m;
}