#!/usr/bin/python

# 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 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.

import os,posix,sys
from subprocess import call
from tempfile import mkstemp

class provider:
    arglist = dict()
    def open(self, provider, header):
        have_provider = False
        self.f = open(provider)
        self.h = open(header,mode='w')
        self.h.write("// Generated by the Systemtap dtrace wrapper\n")
        self.h.write("\n#include <sys/sdt.h>\n\n")
        in_comment = False
        while (True):
            line = self.f.readline()
            if (line == ""):
                break
            if (line.find("/*") >= 0):
                in_comment = True
            if (line.find("*/") >= 0):
                in_comment = False
                continue
            if (in_comment):
                continue
            if (line.find("provider") >= 0):
                tokens = line.split()
                have_provider = True
                self.provider = tokens[1]
            elif (have_provider and line.find("probe ") > 0):
                while (line.find(")") < 0):
                    line += self.f.readline()
                this_probe = line[line.find("probe ")+5:line.find("(")].strip()
                this_probe_canon = self.provider.upper() + "_" + this_probe.replace("__","_").upper()
                args = (line[line.find("(")+1:line.find(")")])
                new_args = ""
                i = 0
                c = 0
                self.arglist[this_probe] = ""
                while (i < len(args)):
                    if (args[i:i+1] == ","):
                        new_args = ('%s%s' % (new_args, args[i]))
                        c += 1
                    else:
                        new_args = new_args + args[i]
                    i += 1
                if (len(new_args) > 0):
                    self.arglist[this_probe] = ('%s arg%d' % (new_args, c))
                if (len(new_args) == 0):
                    self.h.write ('#define %s() STAP_PROBE(provider,%s)\n' % (this_probe_canon, this_probe))
                elif (c == 0):
                    self.h.write ('#define %s(arg1) STAP_PROBE%d(provider,%s,arg1)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 1):
                    self.h.write ('#define %s(arg1,arg2) STAP_PROBE%d(provider,%s,arg1,arg2)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 2):
                    self.h.write ('#define %s(arg1,arg2,arg3) STAP_PROBE%d(provider,%s,arg1,arg2,arg3)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 3):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 4):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 5):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5,arg6) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5,arg6)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 6):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5,arg6,arg7) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5,arg6,arg7)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 7):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 8):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8,arg9) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8,arg9)\n' % (this_probe_canon, c+1, this_probe))
                elif (c == 9):
                    self.h.write ('#define %s(arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8,arg9,arg10) STAP_PROBE%d(provider,%s,arg1,arg2,arg3,arg4,arg5,arg6,arg7,arg8,arg9,arg10)\n' % (this_probe_canon, c+1, this_probe))
                self.h.write ('#define %s_ENABLED() 1\n' % this_probe_canon)

    def get(self, arg):
        print arg
        if (arg in self.arglist):
            return self.arglist[arg]
        else:
            return ""
########################################################################
# main
########################################################################

def usage ():
    print "Usage " + sys.argv[0] + " [-h | -G] -s File.d -o File {Files}"
    sys.exit(1)

def open_file (arg):
    if (len (sys.argv) <= arg):
        return False
    try:
        file = open(sys.argv[arg], 'r')
    except IOError:
        print (sys.argv[arg] + " not found")
        sys.exit(1)
    return file

if (len (sys.argv) < 2):
    usage()

i = 1
build_header = False
build_source = False
filename = ""
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] == "-h"):
        build_header = True
    elif (sys.argv[i] == "-G"):
        build_source = True
    i += 1
if (build_header == False and build_source == False):
    usage()
    sys.exit(1)

if (filename == ""):
    if (s_filename != ""):
	(filename,ext) = os.path.splitext(s_filename)
        filename = os.path.basename(filename)
	if (build_header):
	    filename = filename + ".h"
	elif (build_source):
	    filename = filename + ".o"
    else:
        usage
        sys.exit(1)

if (build_header):
    providers = provider()
    providers.open(s_filename, filename)
elif (build_source):
    (basename,ext) = os.path.splitext(s_filename)
    basename = os.path.basename(basename)
    (d,fn) = mkstemp(suffix=".c",prefix=basename)
    f = open(fn,mode='w')
    f.write("static __dtrace () {}\n")
    f.close()
    call(["gcc", "-fPIC", "-c", fn, "-o", filename], shell=False)
    os.remove(fn)