/*
* Copyright © 2021 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>

#include <funchook.h>

#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
*/

static binaryfunc unicode_cformat_orig;
HOOK_BINARYFUNC(unicode_cformat, "cformat");


#if PY_MAJOR_VERSION < 3
#define PY2

/* python 2 */
/* the hook point for unicode.format & str.format
   are different despite pointing at the same function.
*/
static binaryfunc string_cformat_orig;
HOOK_BINARYFUNC(string_cformat, "cformat");
#else /* python 3 */
#if PY_MINOR_VERSION >= 6
#define FSTRING_HOOK
#if PY_MINOR_VERSION == 6
typedef PyObject ** joinarray_items_t;
#else
/* Apparently they decided to const-correct this function in Py37 */
typedef PyObject *const * joinarray_items_t;
#endif /* PY_MINOR_VERSION == 6 */
PyObject *(*unicode_joinarray_orig)(PyObject *, joinarray_items_t, Py_ssize_t);
#endif /* PY_MINOR_VERSION >= 6 */

static binaryfunc bytes_cformat_orig;
/* bytearray formatting is not implemented in PY2 */
static binaryfunc bytearray_cformat_orig;
HOOK_BINARYFUNC(bytes_cformat, "cformat");
HOOK_BINARYFUNC(bytearray_cformat, "cformat");
#endif /* PY_MAJOR_VERSION < 3 */


#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;

    propagate_result("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 */
int apply_format_patch(funchook_t *funchook) {
#ifdef PY2
    string_cformat_orig = PyString_Format;
    funchook_prep_wrapper(funchook, &string_cformat_orig, string_cformat_new);
#else /* python3 */
    bytes_cformat_orig = PyBytes_Type.tp_as_number->nb_remainder;
    bytearray_cformat_orig = PyByteArray_Type.tp_as_number->nb_remainder;
    funchook_prep_wrapper(funchook, &bytes_cformat_orig, bytes_cformat_new);
    funchook_prep_wrapper(funchook, &bytearray_cformat_orig, bytearray_cformat_new);
#endif /* PY2 */

#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 */

    unicode_cformat_orig = PyUnicode_Format;
    funchook_prep_wrapper(funchook, &unicode_cformat_orig, unicode_cformat_new);

    return 0;
}
