#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-2-Clause
import os
import argparse
import itertools

# Makefile-fuzz-generated.am is created from this template.
MAKEFILE_FUZZ = """# SPDX-License-Identifier: BSD-2-Clause
# Copyright (c) 2018 Intel Corporation
# All rights reserved.

if ENABLE_TCTI_FUZZING
TESTS_FUZZ = %s
%s
endif # ENABLE_TCTI_FUZZING
"""
# Each fuzz target in Makefile-fuzz-generated.am is created from this template.
MAKEFILE_FUZZ_TARGET = """
noinst_PROGRAMS += test/fuzz/%s.fuzz
test_fuzz_%s_fuzz_CPPFLAGS = $(FUZZ_CPPFLAGS)
test_fuzz_%s_fuzz_LDADD    = $(FUZZLDADD)
nodist_test_fuzz_%s_fuzz_SOURCES  = test/fuzz/main-sys.cpp \\
        test/fuzz/%s.fuzz.cpp

DISTCLEANFILES += test/fuzz/%s.fuzz.cpp"""
# Common include definitions needed for fuzzing an SYS call
SYS_TEMPLATE_HEADER = """/* SPDX-License-Identifier: BSD-2-Clause */
/***********************************************************************
 * Copyright (c) 2018, Intel Corporation
 *
 * All rights reserved.
 ***********************************************************************/
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <inttypes.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <poll.h>
#include <stdarg.h>

#include <setjmp.h>

extern "C" {
#include "tss2_mu.h"
#include "tss2_sys.h"
#include "tss2_tcti_device.h"

#include "tss2-tcti/tcti-common.h"
#include "tss2-tcti/tcti-device.h"

#define LOGMODULE fuzz
#include "tss2_tcti.h"
#include "util/log.h"
#include "test.h"
#include "test-options.h"
#include "context-util.h"
#include "tss2-sys/sysapi_util.h"
#include "tcti/tcti-fuzzing.h"
}

extern "C"
int
test_invoke (
        TSS2_SYS_CONTEXT *sysContext)"""
# Template to call a SYS _Complete function which takes no arguments
SYS_COMPLETE_TEMPLATE_NO_ARGS = (
    SYS_TEMPLATE_HEADER
    + """
{
    %s (sysContext);

    return EXIT_SUCCESS;
}
"""
)
# Template to call a SYS _Complete function which takes arguments
SYS_COMPLETE_TEMPLATE_HAS_ARGS = (
    SYS_TEMPLATE_HEADER
    + """
{
    %s

    %s (
        sysContext,
        %s
    );

    return EXIT_SUCCESS;
}
"""
)
# Template to call a SYS _Prepare function
SYS_PREPARE_TEMPLATE_HAS_ARGS = (
    SYS_TEMPLATE_HEADER
    + """
{
    int ret;
    %s

    ret = fuzz_fill (
        sysContext,
        %d,
        %s
    );
    if (ret) {
        return ret;
    }

    %s (
        sysContext,
        %s
    );

    return EXIT_SUCCESS;
}
"""
)


def gen_file(function):
    """
    Generate a cpp file used as the fuzz target given the function definition
    from a header file.
    """
    # Parse the function name from the function definition
    function_name = (
        function.split("\n")[0].replace("TSS2_RC", "").replace("(", "").strip()
    )
    # Parse the function arguments into an array. Do not include sysContext.
    args = [
        arg.strip()
        for arg in function[function.index("(") + 1 : function.index(");")].split(",")
        if not "TSS2_SYS_CONTEXT" in arg
    ]
    # Prepare and Complete functions require different methods of generation.
    # Call the appropriate function to generate a cpp target specific to that
    # type of function.
    if "_Complete" in function_name:
        return gen_complete(function, function_name, args)
    if "_Prepare" in function_name:
        return gen_prepare(function, function_name, args)
    raise NotImplementedError("Unknown function type %r" % (function_name,))


def gen_complete(function, function_name, args):
    """
    Generate the cpp fuzz target for a SYS _Complete call
    """
    if not args:
        # Fill in the no args template. Simple case.
        return function_name, SYS_COMPLETE_TEMPLATE_NO_ARGS % (function_name)
    # Generate the cpp variable definitions.
    arg_definitions = (";\n" + " " * 4).join(
        [arg.replace("*", "") for arg in args]
    ) + ";"
    # Generate the cpp arguments. For arguments that are pointers find replace *
    # with & so that we pass a pointer to the definition which has been
    # allocated on the stack.
    arg_call = (",\n" + " " * 8).join(
        [arg.replace("*", "&").split()[-1] for arg in args]
    )
    # Fill in the template
    return (
        function_name,
        SYS_COMPLETE_TEMPLATE_HAS_ARGS % (arg_definitions, function_name, arg_call),
    )


def gen_prepare(function, function_name, args):
    """
    Generate the cpp fuzz target for a SYS _Prepare call
    """
    if not args:
        return function_name, None
    # Generate the cpp variable definitions. Make sure to initialize to empty
    # structs (works for initializing anything) or cpp compiler will complain.
    arg_definitions = (" = {0};\n" + " " * 4).join(
        [arg.replace("*", "").replace("const", "") for arg in args]
    ) + " = {0};"
    # Generate the cpp arguments. For arguments that are pointers find replace *
    # with & so that we pass a pointer to the definition which has been
    # allocated on the stack.
    arg_call = (",\n" + " " * 8).join(
        [arg.replace("*", "&").split()[-1] for arg in args]
    )
    # Generate the call to fuzz_fill. The call should be the sysContext, double
    # the number of arguments for the _Prepare call, and then for each _Prepare
    # argument pass two to fuzz_fill, the sizeof the _Prepare argument, and a
    # pointer to it.
    fill_fuzz_args = (",\n" + " " * 8).join(
        [
            ("sizeof (%s), &%s" % tuple([arg.replace("*", "").split()[-1]] * 2))
            for arg in args
        ]
    )
    # Fill in the template
    return (
        function_name,
        SYS_PREPARE_TEMPLATE_HAS_ARGS
        % (arg_definitions, len(args) * 2, fill_fuzz_args, function_name, arg_call),
    )


def functions_from_include(header):
    """
    Parse out and yield each function definition from a header file.
    """
    with open(header, "r") as header_fd:
        current_function = ""
        for line in header_fd:
            # Functions we are interested in start with _Complete or _Prepare
            if "_Complete" in line or "_Prepare" in line:
                # Set the current_function to this line
                current_function = line
            elif current_function and ");" in line:
                # When we reach the closing parenthesis yield the function
                yield current_function + line.rstrip()
                current_function = ""
            elif current_function:
                # Add all the arguments to the function
                current_function += line


def gen_files(header):
    # Generate a fuzz target cpp file from each function in the header file
    for current_function in functions_from_include(header):
        function_name, contents = gen_file(current_function)
        # Skip the yield if there is no fuzz target that can be generated
        if contents is None:
            continue
        # Yield the function name and the contents of its generated file
        yield function_name, contents


def main():
    parser = argparse.ArgumentParser(description="Generate libfuzzer for sys")
    parser.add_argument(
        "--header",
        default="include/tss2/tss2_sys.h",
        help="Header file to look in (default include/tss2/tss2_sys.h)",
    )
    args = parser.parse_args()

    functions = dict(gen_files(args.header))
    # Write the generated target to the file for its function name
    for function_name, contents in functions.items():
        filepath = os.path.join("test", "fuzz", function_name + ".fuzz.cpp")
        with open(filepath, "w") as fuzzer_fd:
            fuzzer_fd.write(contents)
    # Fill in the Makefile-fuzz-generated.am template using the function names.
    # Create a list of the compiled fuzz targets
    files = " \\\n    ".join(
        ["test/fuzz/%s.fuzz" % (function) for function in functions]
    )
    # Create the Makefile targets for each generated file
    targets = "\n".join(
        [
            MAKEFILE_FUZZ_TARGET % tuple(list(itertools.chain(([function] * 6))))
            for function in functions
        ]
    )
    # Write out the Makefile-fuzz-generated.am file
    with open("Makefile-fuzz-generated.am", "w") as makefile_fd:
        makefile_fd.write(MAKEFILE_FUZZ % (files, targets))


if __name__ == "__main__":
    main()