#!/usr/bin/env python3
# -*- python -*-
#
# Copyright (C) 2019 Red Hat, Inc
#
# QEMU SystemTap Trace Tool
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.

import argparse
import copy
import os.path
import re
import subprocess
import sys


def probe_prefix(binary):
    dirname, filename = os.path.split(binary)
    return re.sub("-", ".", filename) + ".log"


def which(binary):
    for path in os.environ["PATH"].split(os.pathsep):
        if os.path.exists(os.path.join(path, binary)):
                return os.path.join(path, binary)

    print("Unable to find '%s' in $PATH" % binary)
    sys.exit(1)


def tapset_dir(binary):
    dirname, filename = os.path.split(binary)
    if dirname == '':
        thisfile = which(binary)
    else:
        thisfile = os.path.realpath(binary)
        if not os.path.exists(thisfile):
            print("Unable to find '%s'" % thisfile)
            sys.exit(1)

    basedir = os.path.split(thisfile)[0]
    tapset = os.path.join(basedir, "..", "share", "systemtap", "tapset")
    return os.path.realpath(tapset)


def tapset_env(tapset_dir):
    tenv = copy.copy(os.environ)
    tenv["SYSTEMTAP_TAPSET"] = tapset_dir
    return tenv

def cmd_run(args):
    prefix = probe_prefix(args.binary)
    tapsets = tapset_dir(args.binary)

    if args.verbose:
        print("Using tapset dir '%s' for binary '%s'" % (tapsets, args.binary))

    probes = []
    for probe in args.probes:
        probes.append("probe %s.%s {}" % (prefix, probe))
    if len(probes) == 0:
        print("At least one probe pattern must be specified")
        sys.exit(1)

    script = " ".join(probes)
    if args.verbose:
        print("Compiling script '%s'" % script)
        script = """probe begin { print("Running script, <Ctrl>-c to quit\\n") } """ + script

    # We request an 8MB buffer, since the stap default 1MB buffer
    # can be easily overflowed by frequently firing QEMU traces
    stapargs = ["stap", "-s", "8"]
    if args.pid is not None:
        stapargs.extend(["-x", args.pid])
    stapargs.extend(["-e", script])
    subprocess.call(stapargs, env=tapset_env(tapsets))


def cmd_list(args):
    tapsets = tapset_dir(args.binary)

    if args.verbose:
        print("Using tapset dir '%s' for binary '%s'" % (tapsets, args.binary))

    def print_probes(verbose, name):
        prefix = probe_prefix(args.binary)
        offset = len(prefix) + 1
        script = prefix + "." + name

        if verbose:
            print("Listing probes with name '%s'" % script)
        proc = subprocess.Popen(["stap", "-l", script],
                                stdout=subprocess.PIPE,
                                universal_newlines=True,
                                env=tapset_env(tapsets))
        out, err = proc.communicate()
        if proc.returncode != 0:
            print("No probes found, are the tapsets installed in %s" % tapset_dir(args.binary))
            sys.exit(1)

        for line in out.splitlines():
            if line.startswith(prefix):
                print("%s" % line[offset:])

    if len(args.probes) == 0:
        print_probes(args.verbose, "*")
    else:
        for probe in args.probes:
            print_probes(args.verbose, probe)


def main():
    parser = argparse.ArgumentParser(description="QEMU SystemTap trace tool")
    parser.add_argument("-v", "--verbose", help="Print verbose progress info",
                        action='store_true')

    subparser = parser.add_subparsers(help="commands")
    subparser.required = True
    subparser.dest = "command"

    runparser = subparser.add_parser("run", help="Run a trace session",
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     epilog="""

To watch all trace points on the qemu-system-x86_64 binary:

   %(argv0)s run qemu-system-x86_64

To only watch the trace points matching the qio* and qcrypto* patterns

   %(argv0)s run qemu-system-x86_64 'qio*' 'qcrypto*'
""" % {"argv0": sys.argv[0]})
    runparser.set_defaults(func=cmd_run)
    runparser.add_argument("--pid", "-p", dest="pid",
                           help="Restrict tracing to a specific process ID")
    runparser.add_argument("binary", help="QEMU system or user emulator binary")
    runparser.add_argument("probes", help="Probe names or wildcards",
                           nargs=argparse.REMAINDER)

    listparser = subparser.add_parser("list", help="List probe points",
                                      formatter_class=argparse.RawDescriptionHelpFormatter,
                                      epilog="""

To list all trace points on the qemu-system-x86_64 binary:

   %(argv0)s list qemu-system-x86_64

To only list the trace points matching the qio* and qcrypto* patterns

   %(argv0)s list qemu-system-x86_64 'qio*' 'qcrypto*'
""" % {"argv0": sys.argv[0]})
    listparser.set_defaults(func=cmd_list)
    listparser.add_argument("binary", help="QEMU system or user emulator binary")
    listparser.add_argument("probes", help="Probe names or wildcards",
                            nargs=argparse.REMAINDER)

    args = parser.parse_args()

    args.func(args)
    sys.exit(0)

if __name__ == '__main__':
    main()