#!/usr/bin/python3
# vim: et sta sts=4 sw=4 ts=8

# This handles the systemtap equivalent of
# $(DTRACE) $(DTRACEFLAGS) -G -s $^ -o $@
# $(DTRACE) $(DTRACEFLAGS) -h -s $^ -o $@
# which is a step that builds DTrace provider and probe definitions

# Copyright (C) 2009-2018 Red Hat Inc.
#
# This file is part of systemtap, and is free software.  You can
# redistribute it and/or modify it under the terms of the GNU General
# Public License (GPL); either version 2, or (at your option) any
# later version.

# ignore line too long, missing docstring, method could be a function,
#        too many public methods
# pylint: disable=C0301
# pylint: disable=C0111
# pylint: disable=R0201
# pylint: disable=R0904

import hashlib
import os
import sys
import time
import atexit
from shlex import split
from subprocess import call
try:
    from pyparsing import alphas, cStyleComment, delimitedList, Group, \
        Keyword, lineno, Literal, nestedExpr, nums, oneOf, OneOrMore, \
        Optional, ParseException, ParserElement, restOfLine, restOfLine, \
        Suppress, SkipTo, Word, ZeroOrMore
    HAVE_PYP = True
except ImportError:
    HAVE_PYP = False


# Common file creation methods for pyparsing and string pattern matching

class _HeaderCreator(object):
    def init_semaphores(self, fdesc):
        # dummy declaration just to make the object file non-empty
        fdesc.write("/* Generated by the Systemtap dtrace wrapper */\n\n")
        fdesc.write("static void __dtrace (void) __attribute__((unused));\n")
        fdesc.write("static void __dtrace (void) {}\n")
        fdesc.write("\n#include <sys/sdt.h>\n\n")

    def init_probes(self, fdesc):
        fdesc.write("/* Generated by the Systemtap dtrace wrapper */\n\n")
        fdesc.write("\n#define _SDT_HAS_SEMAPHORES 1\n\n")
        fdesc.write("\n#define STAP_HAS_SEMAPHORES 1 /* deprecated */\n\n")
        fdesc.write("\n#include <sys/sdt.h>\n\n")

    def add_semaphore(self, this_provider, this_probe):
        # NB: unsigned short is fixed in ABI
        semaphores_def = '\n#if defined STAP_SDT_V1\n'
        semaphores_def += '#define %s_%s_semaphore %s_semaphore\n' % \
                          (this_provider, this_probe, this_probe)
        semaphores_def += '#endif\n'
        semaphores_def += '#if defined STAP_SDT_V1 || defined STAP_SDT_V2 \n'
        semaphores_def += "__extension__ unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % \
            (this_provider, this_probe)
        semaphores_def += '#else\n'
        semaphores_def += "__extension__ unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\"))) __attribute__ ((visibility (\"hidden\")));\n" % \
            (this_provider, this_probe)
        semaphores_def += '#endif\n'
        return semaphores_def

    def add_probe(self, this_provider, this_probe, args):
        stap_str = ""
        this_probe_canon = this_provider.upper() + "_" + this_probe.replace("__", "_").upper()
        define_str = "#define %s(" % (this_probe_canon)
        comment_str = "/* %s (" % (this_probe_canon)

        if len(args) == 0:
            stap_str += "DTRACE_PROBE ("
        else:
            stap_str += "DTRACE_PROBE%d (" % len(args)
        stap_str += "%s, %s" % (this_provider, this_probe)
        i = 0
        while i < len(args):
            if i != 0:
                define_str += ", "
                comment_str += ","
            define_str = define_str + "arg%s" % (i + 1)
            stap_str = stap_str + ", arg%s" % (i + 1)
            for argi in args[i]:
                if len(argi) > 0:
                    comment_str += " %s" % argi
            i += 1
        stap_str += ")"
        comment_str += " ) */"
        define_str += ") \\\n"
        probe_def = '%s\n' % (comment_str)
        probe_def += ('#if defined STAP_SDT_V1\n')
        probe_def += ('#define %s_ENABLED() __builtin_expect (%s_semaphore, 0)\n' % \
                      (this_probe_canon, this_probe))
        probe_def += ('#define %s_%s_semaphore %s_semaphore\n' % \
                      (this_provider, this_probe, this_probe))
        probe_def += ('#else\n')
        probe_def += ('#define %s_ENABLED() __builtin_expect (%s_%s_semaphore, 0)\n' % \
                      (this_probe_canon, this_provider, this_probe))
        probe_def += ('#endif\n')
        # NB: unsigned short is fixed in ABI
        probe_def += ("__extension__ extern unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % \
                      (this_provider, this_probe))
        probe_def += (define_str + stap_str + "\n\n")
        return probe_def

# Parse using pyparsing if it is available

class _PypProvider(_HeaderCreator):
    def __init__(self):
        self.ast = []
        self.bnf = None
        self.dtrace_statements = None

    def dtrace_bnf(self):
        self.current_probe = ""
        if self.dtrace_statements is not None:
            return
        ParserElement.setDefaultWhitespaceChars(' \f\r\n\t\v')
        ident = Word(alphas+"_", alphas+nums+"_$")
        probe_ident = Word(alphas+nums+"_$")
        semi = Literal(";").suppress()
        integer = Word( nums )
        lbrace = Literal("{").suppress()
        rbrace = Literal("}").suppress()
        type_name = ident
        varname = ident
        PROBE = Keyword("probe")
        PROVIDER = Keyword("provider")
        array_size = integer | ident  
        varname_spec = varname + Optional("[" + array_size + "]")
        struct_decl = Group(oneOf("struct union") + varname + Suppress(nestedExpr('{','}')) + semi)
        enum_decl = Group("enum" + varname + Suppress(nestedExpr('{','}')) + semi)
        member_decl = Group((Optional(oneOf("struct unsigned const")) + type_name + Optional(Keyword("const")))
                            + Optional(Word("*"), default="") + Optional(varname_spec))
        struct_typedef = Group(Literal("typedef") + Literal("struct") + varname
                               + Suppress(nestedExpr('{','}'))) + Optional(varname) + semi
        typedef = ZeroOrMore("typedef" + (member_decl)) + semi
        decls = OneOrMore(struct_typedef | struct_decl | typedef | enum_decl)
        def memoize_probe(instring, loc, tokens):
            self.current_probe = tokens[0][1]
            self.current_lineno = lineno(loc,instring)
        probe_decl = Group(PROBE + probe_ident + "(" + Optional(Group(delimitedList(member_decl))) + ")" + Optional(Group(Literal(":") + "(" + Optional(Group(delimitedList(member_decl))) + ")")) + Optional(semi))
        probe_decl.setParseAction(memoize_probe)
        probe_decls = OneOrMore(probe_decl)

        provider_decl = (PROVIDER + Optional(ident)
                         + lbrace + Group(probe_decls) + rbrace + Optional(semi))
        dtrace_statement = Group (SkipTo("provider", include=False) + provider_decl)
        self.dtrace_statements = ZeroOrMore(dtrace_statement)

        cplusplus_linecomment = Literal("//") + restOfLine
        cpp_linecomment = Literal("#") + restOfLine

        self.dtrace_statements.ignore(cStyleComment)
        self.dtrace_statements.ignore(cplusplus_linecomment)
        self.dtrace_statements.ignore(cpp_linecomment)

        self.bnf = self.dtrace_statements

    def semaphore_write(self, fdesc):
        semaphores_def = ""
        self.init_semaphores(fdesc)
        for asti in self.ast:
            if len(asti) == 0:
                continue
            # ignore SkipTo token
            if asti[0] != "provider":
                del asti[0]
            if asti[0] == "provider":
                # list of probes
                for prb in asti[2]:
                    semaphores_def += self.add_semaphore(asti[1], prb[1])
        fdesc.write(semaphores_def)


    def probe_write(self, provider, header):
        hdr = open(header, mode='w')
        self.init_probes(hdr)
        self.dtrace_bnf()
        try:
            try:
                self.ast = self.bnf.parseFile(provider, parseAll=True).asList()
            except TypeError:
                # pyparsing-1.5.0 does not support parseAll
                self.ast = self.bnf.parseFile(provider).asList()
        except ParseException:
            err = sys.exc_info()[1]
            if len(self.current_probe):
                print("Warning: %s:%s:%d: syntax error near:\nprobe %s\n" % (sys.argv[0], provider, self.current_lineno, self.current_probe))
            else:
                print("Warning: %s:%s:%d syntax error near:\n%s\n" % (sys.argv[0], provider, err.lineno, err.line))
            raise err
  
        probes_def = ""
        for asti in self.ast:
            if len(asti) == 0:
                continue
            # ignore SkipTo token
            if asti[0] != "provider":
                del asti[0]
            if asti[0] == "provider":
                # list of probes
                for prb in asti[2]:
                    if prb[3] == ')': # No parsed argument list
                        alist = []
                    else:
                        alist = prb[3]
                    probes_def += self.add_probe(asti[1], prb[1], alist)
        hdr.write(probes_def)
        hdr.close()


# Parse using regular expressions if pyparsing is not available

class _ReProvider(_HeaderCreator):
    def __init__(self):
        self.semaphores_def = "\n"
        self.provider = []

    def __semaphore_append(self, this_probe):
        self.semaphores_def += self.add_semaphore(self.provider, this_probe)

    def semaphore_write(self, fdesc):
        self.init_semaphores(fdesc)
        fdesc.write(self.semaphores_def)

    def probe_write(self, provider, header):
        have_provider = False
        fdesc = open(provider)
        hdr = open(header, mode='w')
        self.init_probes(hdr)
        in_comment = False
        probes_def = ""
        while True:
            line = fdesc.readline()
            if line == "":
                break
            if line.find("/*") != -1:
                in_comment = True
            if line.find("*/") != -1:
                in_comment = False
                continue
            if in_comment:
                continue
            if line.find("provider") != -1:
                tokens = line.split()
                have_provider = True
                self.provider = tokens[1]
            elif have_provider and line.find("probe ") != -1:
                while line.find(")") < 0:
                    line += fdesc.readline()
                this_probe = line[line.find("probe ")+5:line.find("(")].strip()
                argstr = (line[line.find("(")+1:line.find(")")])
                arg = ""
                i = 0
                args = []
                self.__semaphore_append(this_probe)
                while i < len(argstr):
                    if argstr[i:i+1] == ",":
                        args.append(arg.split())
                        arg = ""
                    else:
                        arg = arg + argstr[i]
                    i += 1
                if len(arg) > 0:
                    args.append(arg.split())
                probes_def += self.add_probe(self.provider, this_probe, args)
            elif line.find("}") != -1 and have_provider:
                have_provider = False
        hdr.write(probes_def)
        hdr.close()


def mktemp_determ(sources, suffix):
    # for reproducible-builds purposes, use a predictable tmpfile path
    sha = hashlib.sha256()
    for source in sources:
        sha.update(source.encode('utf-8'))
    fname = ".dtrace-temp." + sha.hexdigest()[:8] + suffix
    tries = 0
    while True:
        tries += 1
        if tries > 100: # if file exists due to previous crash or whatever
            raise Exception("cannot create temporary file \""+fname+"\"")
        try:
            wxmode = 'x' if sys.version_info > (3,0) else 'wx'
            fdesc = open(fname, mode=wxmode)
            break
        except FileExistsError:
            time.sleep(0.1) # vague estimate of elapsed time for concurrent identical gcc job
            pass # Try again

    return fdesc, fname


def usage():
    print("Usage " + sys.argv[0] + " [--help] [-h | -G] [-C [-I<Path>]] -s File.d [-o <File>]")


def dtrace_help():
    usage()
    print("Where -h builds a systemtap header file from the .d file")
    print("      -C when used with -h, also run cpp preprocessor")
    print("      -o specifies an explicit output file name,")
    print("         the default for -G is file.o and -h is file.h")
    print("      -I when running cpp pass through this -I include Path")
    print("      -s specifies the name of the .d input file")
    print("      -G builds a stub file.o from file.d,")
    print("         which is required by some packages that use dtrace.")
    sys.exit(1)


########################################################################
# main
########################################################################

def main():
    if len(sys.argv) < 2:
        usage()
        return 1

    global HAVE_PYP
    i = 1
    build_header = False
    build_source = False
    keep_temps = False
    use_cpp = False
    suffix = ""
    filename = ""
    s_filename = ""
    includes = []
    defines = []
    ignore_options = ["-64", "-32", "-fpic", "-fPIC"]
    ignore_options2 = ["-x"] # with parameter

    while i < len(sys.argv):
        if sys.argv[i] == "-o":
            i += 1
            filename = sys.argv[i]
        elif sys.argv[i] == "-s":
            i += 1
            s_filename = sys.argv[i]
        elif sys.argv[i] == "-C":
            use_cpp = True
        elif sys.argv[i].startswith("-D"):
            defines.append(sys.argv[i])
        elif sys.argv[i] == "-h":
            build_header = True
            suffix = ".h"
        elif sys.argv[i].startswith("-I"):
            includes.append(sys.argv[i])
        elif sys.argv[i] == "-G":
            build_source = True
            suffix = ".o"
        elif sys.argv[i] == "-k":
            keep_temps = True
        elif sys.argv[i] == "--no-pyparsing":
            HAVE_PYP = False
        elif sys.argv[i] == "--types":
            print(sys.argv[0] + ": note: obsolete option --types used")
        elif sys.argv[i] in ignore_options:
            pass                # dtrace users sometimes pass these flags
        elif sys.argv[i] in ignore_options2:
            i += 1
            pass                # dtrace users sometimes pass these flags
        elif sys.argv[i] == "--help":
            dtrace_help()
        elif sys.argv[i][0] == "-":
            print(sys.argv[0], "invalid option", sys.argv[i])
            usage()
            return 1
        i += 1
    if not build_header and not build_source:
        usage()
        return 1

    if s_filename != "" and use_cpp:
        (ignore, fname) = mktemp_determ(["use_cpp", s_filename], suffix=".d")
        cpp = os.environ.get("CPP", "cpp")
        retcode = call(split(cpp) + includes + defines + [s_filename, '-o', fname])
        if retcode != 0:
            print("\"cpp includes s_filename\" failed")
            usage()
            return 1
        s_filename = fname
    if filename == "":
        if s_filename != "":
            (filename, ignore) = os.path.splitext(s_filename)
            filename = os.path.basename(filename)
        else:
            usage()
            return 1
    else:
        suffix = ""

    if build_header:
        if HAVE_PYP:
            providers = _PypProvider()
        else:
            providers = _ReProvider()
        while True:
            try:
                providers.probe_write(s_filename, filename + suffix)
                break;
            # complex C declarations can fool the pyparsing grammar.  
            # we could increase the complexity of the grammar
            # instead we fall back to string pattern matching
            except ParseException:
                err = sys.exc_info()[1]
                print("Warning: Proceeding as if --no-pyparsing was given.\n")
                providers = _ReProvider()
    elif build_source:
        if HAVE_PYP:
            providers = _PypProvider()
        else:
            providers = _ReProvider()
        (fdesc, fname) = mktemp_determ(["build_source", s_filename], suffix=".h")
        while True:
            try:
                providers.probe_write(s_filename, fname)
                break;
            except ParseException:
                err = sys.exc_info()[1]
                print("Warning: Proceeding as if --no-pyparsing was given.\n")
                providers = _ReProvider()
        if not keep_temps:
            os.remove(fname)
        else:
            print("header: " + fname)

        (fdesc, fname) = mktemp_determ(["build_source", s_filename, filename], suffix=".c")
        if not keep_temps:
            atexit.register(os.remove, fname) # delete generated source at exit, even if error
        providers.semaphore_write(fdesc)
        fdesc.close()
        cc1 = os.environ.get("CC", "gcc")
        cflags = "-g " + os.environ.get("CFLAGS", "").replace('\\\n', ' ').replace('\\\r',' ')
        # sanitize any embedded \n etc. goo; PR21063
        retcode = call(split(cc1) + defines + includes + split(cflags) +
             ["-fPIC", "-I.", "-I/usr/include", "-c", fname, "-o",
              filename + suffix], shell=False)
        if retcode != 0:
            print("\"gcc " + fname + "\" failed")
            usage()
            return 1
        if keep_temps:
            print("source: " + fname)

    if use_cpp:
        if not keep_temps:
            os.remove(s_filename)
        else:
            print("cpp: " + s_filename)

    return 0

if __name__ == "__main__":
    sys.exit(main())

# Local Variables:
# mode: python
# End: