/*
* 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 <assert.h>
#include <ctype.h>
#include <stddef.h>
#include <stdlib.h>

#include <funchook.h>

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


/*
* Concat:
*
* Combines two strings into a new string
*
* Source: Origin/Self
* Target: Return
* Action: Append
*/


void (*unicode_append_orig)(PyObject **l, PyObject *r);
binaryfunc unicode_concat_orig;
binaryfunc bytes_concat_orig;
binaryfunc bytearray_concat_orig;
binaryfunc bytearray_iconcat_orig;


void unicode_append_new(PyObject **l, PyObject *r) {
    PyObject *orig_l = *l;
    Py_XINCREF(orig_l);

    unicode_append_orig(l, r);

    if (*l == NULL)
        goto cleanup_and_return;

    propagate_concat(orig_l, r, *l);

cleanup_and_return:
    Py_XDECREF(orig_l);
}


PyObject *unicode_concat_new(PyObject *l, PyObject *r) {
    PyObject *result = unicode_concat_orig(l, r);

    if (result == NULL)
        return result;

    propagate_concat(l, r, result);

    return result;
}


PyObject *bytes_concat_new(PyObject *l, PyObject *r) {
    PyObject *result = bytes_concat_orig(l, r);

    if (result == NULL)
        return result;

    propagate_concat(l, r, result);

    return result;
}


PyObject *bytearray_concat_new(PyObject *l, PyObject *r) {
    PyObject *result = bytearray_concat_orig(l, r);

    if (result == NULL)
        return result;

    propagate_concat(l, r, result);

    return result;
}


/* in-place concatenation */
PyObject *bytearray_iconcat_new(PyObject *self, PyObject *other) {
    PyObject *result = bytearray_iconcat_orig(self, other);

    if (result == NULL)
        return result;

    propagate_concat(self, other, result);

    return result;
}


int apply_cat_patch(funchook_t *funchook) {
    unicode_append_orig = PyUnicode_Append;
    unicode_concat_orig = PyUnicode_Concat;
    bytes_concat_orig = PyBytes_Type.tp_as_sequence->sq_concat;
    bytearray_concat_orig = PyByteArray_Concat;
    bytearray_iconcat_orig = PyByteArray_Type.tp_as_sequence->sq_inplace_concat;

    funchook_prep_wrapper(funchook, &unicode_append_orig, unicode_append_new);
    funchook_prep_wrapper(funchook, &unicode_concat_orig, unicode_concat_new);
    funchook_prep_wrapper(funchook, &bytes_concat_orig, bytes_concat_new);
    funchook_prep_wrapper(funchook, &bytearray_concat_orig, bytearray_concat_new);
    funchook_prep_wrapper(funchook, &bytearray_iconcat_orig, bytearray_iconcat_new);

    return 0;
}
