#-------------------------------------------------------------------------------
# Converted from elftools example: elf_low_high_api.py
#
# Eli Bendersky (eliben@gmail.com)
# This code is in the public domain
#-------------------------------------------------------------------------------
#
# Code in this file counts libcw function symbols in libcw's .so file
# (function's whose names start with "cw_"). The .so file must be
# generated by "make check" command to have more library functions
# visible to the tool.
# The code also discovers and counts functions in libcw's test code,
# marked with "LIBCW_TEST_FUT()" macro - those functions are tested by
# the test code.
# Then the code counts percentage of tested functions in relation to
# all functions from library. Thanks to the calculated value we have
# additional information about test coverage.
#
# python3 ./this_script --so <path to libcw_test.so file> --code <path to dir with test code testing libcw functions>
#
# Kamil Ignacak (not a python master, obviously), 2020
#


import sys
import os

# If pyelftools is not installed, the example can also run from the root or
# examples/ dir of the source distribution.
sys.path[0:0] = ['.', '..']

from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection




def find_libcw_functions_in_so(filename):
    result = set()
    with open(filename, 'rb') as f:
        result = section_info_highlevel(f)
    return result




def section_info_highlevel(stream):
    result = set()
    elffile = ELFFile(stream)

    # Just use the public methods of ELFFile to get what we need
    # Note that section names are strings.
    print("  {} sections".format(elffile.num_sections()))
    section = elffile.get_section_by_name('.symtab')

    if not section:
        print('  No symbol table found. Perhaps this ELF has been stripped?')
        return result

    # A section type is in its header, but the name was decoded and placed in
    # a public attribute.
    print('  Section name: %s, type: %s' %(section.name, section['sh_type']))

    # But there's more... If this section is a symbol table section (which is
    # the case in the sample ELF file that comes with the examples), we can
    # get some more information about it.
    if not isinstance(section, SymbolTableSection):
        return result

    num_symbols = section.num_symbols()
    print("  It's a symbol section with %s symbols" % num_symbols)
    for i in range(0, num_symbols):
        symbol_type = section.get_symbol(i)['st_info']['type'];
        if symbol_type != 'STT_FUNC':
            continue
        symbol_name = section.get_symbol(i).name
        if symbol_name[0:3] != "cw_":
            continue
        if "localalias" in symbol_name: # e.g. cw_send_representation_partial.localalias.2; TODO: why do we have symbols like this one?
            continue
        result.add(symbol_name)

    return result




# Find functions put in 'fut_tag' in given file with test code. This may look
# like this in source code:
#
# cw_ret_t cwret = LIBCW_TEST_FUT(cw_gen_remove_last_character)(gen);
#
# Here LIBCW_TEST_FUT is the tag, and cw_gen_remove_las6t_character is the
# searched function name. Notice that function's parens around list of
# arguments are outside of fut tag.
def find_fut_tagged_functions_in_test_file(file_handle, fut_tag):
    result = set()
    while True:
        line = file_handle.readline()
        if not line:
            break

        idx_start = line.find(fut_tag)
        if -1 == idx_start:
            continue
        idx_start += len(fut_tag + "(") # tag + opening paren: "LIBCW_TEST_FUT("

        # Find closing paren of FUT macro; should cover both usages of
        # FUT macro: on a name of called function, and on a name of
        # function pointer variable.
        idx_end = line.find(")", idx_start)
        if -1 == idx_end:
            continue

        function = line[idx_start:idx_end]
        if not function:
            continue

        result.add(function)

    return result




def find_fut_tagged_functions_in_dir(directory, fut_tag):
    result = set()
    return find_fut_tagged_functions_in_dir2(result, directory, fut_tag)




def find_fut_tagged_functions_in_dir2(result, directory, fut_tag):
    with os.scandir(directory) as it:
        for entry in it:
            if entry.is_file():
                if not entry.name.endswith(".c"):
                    continue
                with open(directory + "/" + entry.name) as file:
                    result |= find_fut_tagged_functions_in_test_file(file, fut_tag) # merge two sets
            elif entry.is_dir():
                # Recurse into dir.
                result |= find_fut_tagged_functions_in_dir2(result, entry.name, fut_tag)

    return result




# Print a set with function names: global functions first, internal
# functions second.
def print_functions_set(header, functions):
    print("\n" + header)

    for fun in sorted(functions):
        if not fun.endswith("_internal"):
            print("    " + fun)

    print("")

    for fun in sorted(functions):
        if fun.endswith("_internal"):
            print("    " + fun)

    print("")




def print_function_row(function_counter, function_name, is_in_so, is_in_fut):
    print('{:>3}   {:<55}'.format(function_counter, function_name), end='')

    if is_in_so:
        print("       ", end='')
    else:
        print("    so-", end='')

    if is_in_fut:
        print("        ")
    else:
        print("    fut-")




def print_stats(n_tested, n_all):
    print("Functions under test: {} out of {} ({:.2f}%)".format(n_tested, n_all, 100.0 *  n_tested/n_all))




def print_help(script_name):
    print("Count how many libcw functions are tested")
    print("Call it like this:\n")
    print("python3 {}\n"
          "    or\n"
          "python3 {} --so <path to libcw so file> --code <path to dir with test code>".format(script_name, script_name))




if __name__ == '__main__':
    so_symbols = set()
    fut_tagged_functions = set()
    so_location = ""
    test_code_location = ""

    if len(sys.argv) == 1:
        so_location = "../.libs/libcw.so"
        test_code_location = "./"
    elif len(sys.argv) == 5:
        if sys.argv[1] == '--so' and sys.argv[3] == '--code':
            so_location = sys.argv[2]
            test_code_location = sys.argv[4]
        else:
            print_help(sys.argv[0])
            sys.exit(1)
    else:
        print_help(sys.argv[0])
        sys.exit(1)

    so_symbols = find_libcw_functions_in_so(so_location)
    fut_tagged_functions = find_fut_tagged_functions_in_dir(test_code_location, "LIBCW_TEST_FUT")

    #print_functions_set("All libcw.so function symbols:", so_symbols)
    #print_functions_set("\nlibcw FUT-tagged functions (functions under test):", fut_tagged_functions)
    #print_functions_set("libcw functions not tagged with FUT yet:", so_symbols - fut_tagged_functions)
    #print("\n")


    all_functions = so_symbols | fut_tagged_functions
    n_tested = len(fut_tagged_functions)
    n_all = len(all_functions)


    # Check if all tested symbols were found of list of symbols in
    # library:
    #if not fut_tagged_functions.issubset(so_symbols):
    #    print("Error: set of FUT functions is not a subset of all so symbols")
    #    # First build set of functions that are not common, and then preserve only tested functions.
    #    not_in_so = (so_symbols ^ fut_tagged_functions) & fut_tagged_functions
    #    print_functions_set("FUT functions not in .so:", not_in_so)

    print_stats(n_tested, n_all)

    print("")
    print("'fut-' - function not tested (not marked by 'function under test' macro")
    print("'so-'  - function not found in .so library");
    print("")

    unexpected_functions = 0;

    print("Public functions:\n")
    fun_counter = 1
    for fun in sorted(all_functions):
        if fun.count(".") > 0:
            unexpected_functions += 1
            continue
        if not fun.endswith("_internal"):
            is_in_so = fun in so_symbols
            is_in_fut = fun in fut_tagged_functions
            print_function_row(fun_counter, fun, is_in_so, is_in_fut)
            fun_counter += 1

    print("\nPrivate functions:\n")
    for fun in sorted(all_functions):
        if fun.count(".") > 0:
            unexpected_functions += 1
            continue
        if fun.endswith("_internal"):
            is_in_so = fun in so_symbols
            is_in_fut = fun in fut_tagged_functions
            print_function_row(fun_counter, fun, is_in_so, is_in_fut)
            fun_counter += 1

    if unexpected_functions > 0:
        # Functions that should not have been found by this tool. Examples:
        # cw_gen_value_tracking_set_value_internal.constprop.0
        # cw_key_ik_set_value_internal.isra.0
        # cw_timer_run_internal.part.0
        print("\nUnexpected functions:\n")
        for fun in sorted(all_functions):
            if fun.count(".") > 0:
                is_in_so = fun in so_symbols
                is_in_fut = fun in fut_tagged_functions
                print_function_row(fun_counter, fun, is_in_so, is_in_fut)
                fun_counter += 1


    # Print the stats again after a long long list of functions has been
    # printed.
    print("")
    print_stats(n_tested, n_all)
