/*
 * Copyright © 2023 Contrast Security, Inc.
 * See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 */
/* Python requires its own header to always be included first */
#define PY_SSIZE_T_CLEAN
#include <Python.h>

#ifndef NO_FUNCHOOK
#include <funchook.h>
#endif

#include <contrast/assess/propagate.h>
#include <contrast/assess/scope.h>
#include <contrast/assess/utils.h>

/*
 * Format:
 *
 * Formats a string with a list of arguments
 *
 * Source: Origin/Self, All args and kwargs
 * Target: Return
 * Action: Splat
 */
#ifndef NO_FUNCHOOK
static binaryfunc unicode_cformat_orig;
HOOK_BINARYFUNC(unicode_cformat);

#if PY_MINOR_VERSION > 6
#define FSTRING_HOOK
/* Apparently they decided to const-correct this function in Py37 */
typedef PyObject *const *joinarray_items_t;
PyObject *(*unicode_joinarray_orig)(PyObject *, joinarray_items_t, Py_ssize_t);
#endif /* PY_MINOR_VERSION > 6 */
#endif /* ifndef NO_FUNCHOOK */

static binaryfunc bytes_cformat_orig;

static binaryfunc bytearray_cformat_orig;
HOOK_BINARYFUNC(bytes_cformat);
HOOK_BINARYFUNC(bytearray_cformat);

#ifdef FSTRING_HOOK
/*
 * It turns out that f-strings basically are converted into string joins under
 * the hood. Each of the format arguments are evaluated and then joined with
 * the literal parts of the original string. We simply hook the function that
 * is called by the interpreter and call our join propagator.
 */
static PyObject *unicode_joinarray_new(
    PyObject *sep, joinarray_items_t items, Py_ssize_t seqlen) {

    PyObject *result = unicode_joinarray_orig(sep, items, seqlen);
    PyObject *item_list = pack_args_tuple(items, seqlen);
    PyObject *args = PyTuple_Pack(1, item_list);

    if (result == NULL)
        goto cleanup_and_exit;

    call_string_propagator("propagate_unicode_fstring", sep, result, args, NULL);

cleanup_and_exit:
    Py_XDECREF(args);
    Py_XDECREF(item_list);
    return result;
}
#endif /* FSTRING_HOOK */

/* Hooks for the .format method are autogenerated and not included here */
#ifdef NO_FUNCHOOK
void apply_format_patch() {
#else
int apply_format_patch(funchook_t *funchook) {
#endif /* ifdef NO_FUNCHOOK */
       /* TODO: PYT-2710 Investigate removing legacy cformat hooks in favor of modulo
        * rewrite */
    bytes_cformat_orig = PyBytes_Type.tp_as_number->nb_remainder;
    bytearray_cformat_orig = PyByteArray_Type.tp_as_number->nb_remainder;
    PyBytes_Type.tp_as_number->nb_remainder = (void *)bytes_cformat_new;
    PyByteArray_Type.tp_as_number->nb_remainder = (void *)bytearray_cformat_new;

#ifndef NO_FUNCHOOK
#ifdef FSTRING_HOOK
    /* NOTE: this function is available but not part of the stable API */
    unicode_joinarray_orig = _PyUnicode_JoinArray;
    funchook_prep_wrapper(funchook, &unicode_joinarray_orig, unicode_joinarray_new);
#endif /* FSTRING_HOOK */
       /* TODO: PYT-2710 Investigate removing legacy cformat hooks in favor of modulo
        * rewrite */
    unicode_cformat_orig = PyUnicode_Format;
    funchook_prep_wrapper(funchook, &unicode_cformat_orig, unicode_cformat_new);

    return 0;
#endif /* ifndef NO_FUNCHOOK */
}

void reverse_format_patch() {
    PyBytes_Type.tp_as_number->nb_remainder = (void *)bytes_cformat_orig;
    PyByteArray_Type.tp_as_number->nb_remainder = (void *)bytearray_cformat_orig;
}
