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

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

typedef PyObject *(*fastcall_func)(PyObject *, PyObject *const *, Py_ssize_t);
typedef PyObject *(*fastcall_kwargs_func)(PyObject *, PyObject *const *, Py_ssize_t, PyObject *);

#define UNICODE_ENCODE_OFFSET 0
#define UNICODE_REPLACE_OFFSET 1
#define UNICODE_SPLIT_OFFSET 2
#define UNICODE_RSPLIT_OFFSET 3
#define UNICODE_JOIN_OFFSET 4
#define UNICODE_CAPITALIZE_OFFSET 5
#define UNICODE_CASEFOLD_OFFSET 6
#define UNICODE_TITLE_OFFSET 7
#define UNICODE_CENTER_OFFSET 8
#define UNICODE_EXPANDTABS_OFFSET 10
#define UNICODE_LJUST_OFFSET 14
#define UNICODE_LOWER_OFFSET 15
#define UNICODE_LSTRIP_OFFSET 16
#define UNICODE_RJUST_OFFSET 19
#define UNICODE_RSTRIP_OFFSET 20
#define UNICODE_SPLITLINES_OFFSET 22
#define UNICODE_STRIP_OFFSET 23
#define UNICODE_SWAPCASE_OFFSET 24
#define UNICODE_TRANSLATE_OFFSET 25
#define UNICODE_UPPER_OFFSET 26
#define UNICODE_ZFILL_OFFSET 43
#define UNICODE_FORMATMAP_OFFSET 45
#define UNICODE_REMOVEPREFIX_OFFSET 29
#define UNICODE_REMOVESUFFIX_OFFSET 30


fastcall_kwargs_func unicode_encode_orig;
fastcall_func unicode_replace_orig;
fastcall_kwargs_func unicode_split_orig;
fastcall_kwargs_func unicode_rsplit_orig;
binaryfunc unicode_join_orig;
unaryfunc unicode_capitalize_orig;
unaryfunc unicode_casefold_orig;
unaryfunc unicode_title_orig;
fastcall_func unicode_center_orig;
fastcall_kwargs_func unicode_expandtabs_orig;
fastcall_func unicode_ljust_orig;
unaryfunc unicode_lower_orig;
fastcall_func unicode_lstrip_orig;
fastcall_func unicode_rjust_orig;
fastcall_func unicode_rstrip_orig;
fastcall_kwargs_func unicode_splitlines_orig;
fastcall_func unicode_strip_orig;
unaryfunc unicode_swapcase_orig;
binaryfunc unicode_translate_orig;
unaryfunc unicode_upper_orig;
binaryfunc unicode_zfill_orig;
binaryfunc unicode_formatmap_orig;
binaryfunc unicode_removeprefix_orig;
binaryfunc unicode_removesuffix_orig;
newfunc unicode_new_orig;


HOOK_TERNARY_FASTCALL(unicode_encode, "encode");
PyObject *unicode_replace_new(PyObject *self, PyObject *const *args, Py_ssize_t nargs) {
    /* In Py37 the replace method type moved to METH_FASTARGS. This means that
     * instead of args being passed as a tuple, they are passed as a C array
     * that contains PyObjects. We need to check whether there is the number of
     * args that we expect, and whether the arg we care about is not NULL.
     * Specifically, we want args[1] since it represents the "new" string in
     * the replacement.
     */
    PyObject *hook_args = pack_args_tuple(args, nargs);
    PyObject *result = unicode_replace_orig(self, args, nargs);

    if (result == NULL || nargs < 2 || args[1] == NULL)
        goto cleanup_and_exit;

    if (result == self)
        goto cleanup_and_exit;

    propagate_result("replace", self, result, hook_args, NULL);

cleanup_and_exit:
    Py_XDECREF(hook_args);
    return result;

}

PyObject *unicode_split_new(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) {
    PyObject *result = unicode_split_orig(self, args, nargs, kwnames);
    PyObject *args_tuple = pack_args_tuple(args, nargs);
    PyObject *kwargs = pack_kwargs_dict(args, nargs, kwnames);

    if (result == NULL || PySequence_Length(result) == 1)
        goto cleanup_and_exit;

    propagate_result("split", (PyObject *)self, result, args_tuple, kwargs);

cleanup_and_exit:
    Py_XDECREF(args_tuple);
    Py_XDECREF(kwargs);
    return result;

}

PyObject *unicode_rsplit_new(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) {
    PyObject *result = unicode_rsplit_orig(self, args, nargs, kwnames);
    PyObject *args_tuple = pack_args_tuple(args, nargs);
    PyObject *kwargs = pack_kwargs_dict(args, nargs, kwnames);

    if (result == NULL || PySequence_Length(result) == 1)
        goto cleanup_and_exit;

    propagate_result("rsplit", (PyObject *)self, result, args_tuple, kwargs);

cleanup_and_exit:
    Py_XDECREF(args_tuple);
    Py_XDECREF(kwargs);
    return result;

}

PyObject *unicode_join_new(PyObject *self, PyObject *args) {
    PyObject *list = PySequence_List(args);

    /* If this fails for any reason, just call the original function and get
     * out of here.
     */
    if (list == NULL) {
        PyErr_Clear();
        return unicode_join_orig((PyObject *)self, args);
    }

    /* In Py36+ we also hook an internal function that is called by this
     * function in order to propagate fstring formatting. We still want to have
     * a separate hook for join so that the events are reported differently.
     * This means that we need to go into scope when calling the original
     * function here so that we don't propagate twice.
     */
    enter_propagation_scope();
    PyObject *result = unicode_join_orig((PyObject *)self, list);
    exit_propagation_scope();

    PyObject *prop_args = PyTuple_Pack(1, list);

    if (prop_args == NULL || result == NULL)
        goto cleanup_and_exit;

    propagate_result("join", (PyObject *)self, result, prop_args, NULL);

cleanup_and_exit:
    Py_XDECREF(list);
    Py_XDECREF(prop_args);
    return result;

}

HOOK_UNARYFUNC(unicode_capitalize, "capitalize");
HOOK_UNARYFUNC(unicode_casefold, "casefold");
HOOK_UNARYFUNC(unicode_title, "title");
HOOK_FASTCALL(unicode_center, "center");
HOOK_TERNARY_FASTCALL(unicode_expandtabs, "expandtabs");
HOOK_FASTCALL(unicode_ljust, "ljust");
HOOK_UNARYFUNC(unicode_lower, "lower");
HOOK_FASTCALL(unicode_lstrip, "lstrip");
HOOK_FASTCALL(unicode_rjust, "rjust");
HOOK_FASTCALL(unicode_rstrip, "rstrip");
HOOK_TERNARY_FASTCALL(unicode_splitlines, "splitlines");
HOOK_FASTCALL(unicode_strip, "strip");
HOOK_UNARYFUNC(unicode_swapcase, "swapcase");
HOOK_BINARYFUNC(unicode_translate, "translate");
HOOK_UNARYFUNC(unicode_upper, "upper");
HOOK_BINARYFUNC(unicode_zfill, "zfill");
PyObject *unicode_formatmap_new(PyObject *self, PyObject *args) {
    PyObject *result = unicode_formatmap_orig(self, args);

    if (result == NULL)
        return result;

    propagate_result("formatmap", self, result, NULL, args);

    return result;

}

HOOK_BINARYFUNC(unicode_removeprefix, "removeprefix");
HOOK_BINARYFUNC(unicode_removesuffix, "removesuffix");
PyObject *unicode_new_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
    PyObject *result = unicode_new_orig(type, args, kwds);

    if (result == NULL)
        return result;

    propagate_result("CAST", NULL, result, args, kwds);

    return result;

}




int apply_unicode_patches(funchook_t *funchook) {

    ADD_METHOD_HOOK(PyUnicode_Type, unicode_encode, UNICODE_ENCODE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_replace, UNICODE_REPLACE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_split, UNICODE_SPLIT_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_rsplit, UNICODE_RSPLIT_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_join, UNICODE_JOIN_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_capitalize, UNICODE_CAPITALIZE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_casefold, UNICODE_CASEFOLD_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_title, UNICODE_TITLE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_center, UNICODE_CENTER_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_expandtabs, UNICODE_EXPANDTABS_OFFSET);
    /* unicode.partition hook is not included because it is buggy */
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_ljust, UNICODE_LJUST_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_lower, UNICODE_LOWER_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_lstrip, UNICODE_LSTRIP_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_rjust, UNICODE_RJUST_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_rstrip, UNICODE_RSTRIP_OFFSET);
    /* unicode.rpartition hook is not included because it is buggy */
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_splitlines, UNICODE_SPLITLINES_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_strip, UNICODE_STRIP_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_swapcase, UNICODE_SWAPCASE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_translate, UNICODE_TRANSLATE_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_upper, UNICODE_UPPER_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_zfill, UNICODE_ZFILL_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_formatmap, UNICODE_FORMATMAP_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_removeprefix, UNICODE_REMOVEPREFIX_OFFSET);
    ADD_METHOD_HOOK(PyUnicode_Type, unicode_removesuffix, UNICODE_REMOVESUFFIX_OFFSET);
    ADD_NEWFUNC_HOOK(PyUnicode_Type, unicode_new);

    /* Indicates success */
    return 0;
}

