#!/bin/bash

# $1 - install/uninstall
# $2 - PID/unique name
# $3 - RULE name
# $4 - class
# $5 - method
# $6 - number of args
# $7 - entry/exit/line
# $8 - backtrace trigger

exec 1>&2  # redirect byteman/etc. tracing output to stderr, for easier filtering

if [[ $# -gt 8 || $# -lt 7 ]]; then
    echo "need seven or eight arguments"
    exit 1
fi

arg_command=$1
# resolve the overloaded parameter; PR21020
if [ $arg_command = "install31" ]; then
    mode=install
    stap="31"
elif [ $arg_command = "uninstall31" ]; then
    mode=uninstall
    stap="31"
elif [ $arg_command = "install" ]; then
    mode=install
    stap=""
elif [ $arg_command = "uninstall" ]; then
    mode=uninstall
else
    exit 1
fi
arg_jvmpid=$2
arg_rulename=$3
arg_class=$4
arg_method=$5
arg_argcount=$6
arg_probetype=$7

if [ $# -eq 7 ]; then
    arg_backtrace=0
else
    arg_backtrace=$8
fi

SYSTEMTAP_DIR=${SYSTEMTAP_DIR-$HOME/.systemtap}
BYTEMAN_HOME=${BYTEMAN_HOME-/usr/share/java/byteman}
JAVA_HOME=${JAVA_HOME-/usr/lib/jvm/java}
BYTEMAN_INSTALL_OPTS=${BYTEMAN_INSTALL_OPTS--Dorg.jboss.byteman.transform.all=true}
SYSTEMTAP_VERBOSE=${SYSTEMTAP_VERBOSE-0}

if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
    BYTEMAN_INSTALL_OPTS="$BYTEMAN_INSTALL_OPTS -Dorg.jboss.byteman.verbose"
else
    exec >/dev/null
    # NB: preserve stderr
fi

# the byteman and byteman-submit jars should be in ${BYTEMAN_HOME}/lib
BYTEMAN_JAR=${BYTEMAN_HOME}/byteman.jar
if [ ! -r ${BYTEMAN_JAR} ]; then
    exec 1>&2
    echo "Missing $BYTEMAN_JAR"
    exit 1
fi

BYTEMAN_SUBMIT_JAR=${BYTEMAN_HOME}/byteman-submit.jar
if [ ! -r ${BYTEMAN_SUBMIT_JAR} ]; then
    exec 1>&2
    echo "Missing $BYTEMAN_SUBMIT_JAR"
    exit 1
fi

BYTEMAN_INSTALL_JAR=${BYTEMAN_HOME}/byteman-install.jar
if [ ! -r ${BYTEMAN_INSTALL_JAR} ]; then
    exec 1>&2
    echo "Missing $BYTEMAN_INSTALL_JAR"
    exit 1
fi


TOOLS_JAR=${JAVA_HOME}/lib/tools.jar
if [ ! -f ${TOOLS_JAR} ]; then
    exec 1>&2
    echo "Missing $TOOLS_JAR"
    exit 1
fi

# resolve $*prefix fully
prefix=/usr
exec_prefix=/usr
pkglibexecdir=/usr/libexec/systemtap
pkglibexecdir=`eval echo $pkglibexecdir`
pkglibexecdir=`eval echo $pkglibexecdir`

HELPERSDT_JAR=${pkglibexecdir}/HelperSDT.jar
if [ ! -f ${HELPERSDT_JAR} ]; then
    exec 1>&2
    echo "Missing $HELPERSDT_JAR"
    exit 1
fi


# The JVM that invokes byteman will get classpath/etc. settings from
# this shell script to look directly under our $prefix.  However, for
# the target JVM, the HelperSDT* stuff needs to be installed under
# that JVM's paths.

num=`ls -1 ${JAVA_HOME}/jre/lib/ext/HelperSDT.jar ${JAVA_HOME}/jre/lib/*/libHelperSDT_*.so 2>/dev/null | wc -l`
if [ $num -lt 2 ]; then
    exec 1>&2
    echo "Missing HelperSDT JNI class/shared library"
    echo "Install them like this, as root:"
    echo ""
    echo "for so in ${pkglibexecdir}/libHelperSDT_*.so; do"
    echo '  arch=`basename $so | cut -f2 -d_ | cut -f1 -d.`'
    echo "  ln -sf ${pkglibexecdir}/libHelperSDT_"'${arch}'".so ${JAVA_HOME}/jre/lib/"'${arch}'"/"
    echo "done"
    echo "ln -sf ${pkglibexecdir}/HelperSDT.jar ${JAVA_HOME}/jre/lib/ext/"
#   exit 1
fi

flagdir="$SYSTEMTAP_DIR/java"
mkdir -p $flagdir

# Find our target jvm pid.  Due to the possibility of our
# target jvm pid being passed as a string, we need to allow
# for the possiblity that more than one pid may match the
# target jvm pid.  If this is the case, we need to have a
# nested call to stapbm with the actual pid of the jvm pid

if ! [[ $arg_jvmpid =~ ^[0-9]+$ ]]; then
    target_pid=`jps -l | grep $arg_jvmpid | cut -f1 -d" "`
    for target in $target_pid; do
	$0 $arg_command $target $arg_rulename $arg_class "$arg_method" $arg_argcount $arg_probetype $arg_backtrace
    done;
    exit 0
else
    target_pid=$arg_jvmpid
fi

# Our target jvm may not have the byteman agent installed yet.  Let's do
# that first.  We use a signal file in $flagdir to show that the
# JVM is ready for further bytemanning without a prior setup step,
# and include in it the designated byteman agent listening-port number.
#
byteman_installed_portfile=$flagdir/`hostname`-${target_pid}-bm

exec 200>>$byteman_installed_portfile # open/create lock file
flock -x 200  # exclusive-lock it

if [ -s $byteman_installed_portfile ]; then
    bmport=`cat $byteman_installed_portfile`

    if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
        echo "Byteman agent reused for java pid $target_pid, port $bmport"
    fi


    # XXX: liveness-check the port; bmsubmit with no argument just lists current rules
    # if fails, delete the _portfile and retry everything
else
    # bmport=9091
    bmport=`expr 9090 + $RANDOM % 10000`
    existing=`ss -atn | awk '{print $4}' | grep ':'$bmport'$'`
    if [ "x$existing" != "x" ]; then
        echo "Byteman port $bmport already in use, retrying."
        exec "$@"
    fi

# There are two ways to invoke and run byteman operations with the jvm's we're interested
# in, we can alter the startup arguments to include a -javaagent parameter, or use
# byteman and its use of VMAttach libraries, for our case it always makes sense to use
# byteman classes directly and avoid -javaagent

    if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
	echo java -classpath ${BYTEMAN_INSTALL_JAR}:${BYTEMAN_JAR}:${TOOLS_JAR}:${HELPERSDT_JAR} org.jboss.byteman.agent.install.Install -b -p $bmport $BYTEMAN_INSTALL_OPTS $target_pid
    fi
    java -classpath ${BYTEMAN_INSTALL_JAR}:${BYTEMAN_JAR}:${TOOLS_JAR}:${HELPERSDT_JAR} org.jboss.byteman.agent.install.Install -b -p $bmport $BYTEMAN_INSTALL_OPTS $target_pid
    if [ $? -ne 0 ]; then
        echo "Byteman agent failed to install for java pid $target_pid, port $bmport"
        exit 1
    fi
    echo $bmport > $byteman_installed_portfile

    if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
        echo "Byteman agent installed for java pid $target_pid, port $bmport"
    fi
    # XXX: Erase file to keep it from sticking around indefinitely,
    # in case process ends, machine reboots, pid gets reused
    # XXX: consider explicit notification to stapbm via process("java").begin/end ?
    # ... or else: liveness-check below
fi
exec 200>&-   # close file & release flock


function echo_bytemanrule()
{
    echo "RULE $arg_rulename"
    echo "CLASS $arg_class"
    echo "METHOD $arg_method"
    echo "HELPER org.systemtap.byteman.helper.HelperSDT"
    case "$arg_probetype" in
        entry)
	    echo "AT ENTRY"
	    ;;
        exi*)
	    echo "AT RETURN"
	    ;;
        *)
	    echo "AT LINE $arg_probetype"
	    ;;
    esac
    echo "IF TRUE"
    if [ "$arg_backtrace" == "1" ]; then
	echo 'DO STAP_BACKTRACE("'$arg_rulename'");'
    else
	echo -n 'DO '
    fi
    case "$arg_argcount" in
        # For PR21010, we invoke another java<->stap ABI
        0) echo -n 'METHOD_STAP'$stap'_PROBE0("'$arg_rulename'")' ;;
        1) echo -n 'METHOD_STAP'$stap'_PROBE1("'$arg_rulename'", $1)' ;;
        2) echo -n 'METHOD_STAP'$stap'_PROBE2("'$arg_rulename'", $1, $2)' ;;
        3) echo -n 'METHOD_STAP'$stap'_PROBE3("'$arg_rulename'", $1, $2, $3)' ;;
        4) echo -n 'METHOD_STAP'$stap'_PROBE4("'$arg_rulename'", $1, $2, $3, $4)' ;;
        5) echo -n 'METHOD_STAP'$stap'_PROBE5("'$arg_rulename'", $1, $2, $3, $4, $5)' ;;
        6) echo -n 'METHOD_STAP'$stap'_PROBE6("'$arg_rulename'", $1, $2, $3, $4, $5, $6)' ;;
        7) echo -n 'METHOD_STAP'$stap'_PROBE7("'$arg_rulename'", $1, $2, $3, $4, $5, $6, $7)' ;;
        8) echo -n 'METHOD_STAP'$stap'_PROBE8("'$arg_rulename'", $1, $2, $3, $4, $5, $6, $7, $8)' ;;
        9) echo -n 'METHOD_STAP'$stap'_PROBE9("'$arg_rulename'", $1, $2, $3, $4, $5, $6, $7, $8, $9)' ;;
	10) echo -n 'METHOD_STAP'$stap'_PROBE10("'$arg_rulename'", $1, $2, $3, $4, $5, $6, $7, $8, $9, $10)' ;;
        *) echo 'bad arg-count'; exit 1 ;;
    esac
    if [ "$arg_backtrace" == "1" ]; then
	echo ';'
	echo 'METHOD_BT_DELETE("'$arg_rulename'")'
    else
	echo ''
    fi
    echo "ENDRULE"
}


# Generate the byteman rule file on-the-fly
btmfile=$flagdir/`hostname`-$$.btm
echo_bytemanrule > $btmfile
trap 'rm -f $btmfile' 0 1 2 3 4 5 9 15

if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
    echo "Byteman rule file:"
    cat $btmfile
fi

if [ $mode = "uninstall" ]; then
    bmcmd=-u
else
    bmcmd=-l
fi

if [ "$SYSTEMTAP_VERBOSE" != "0" ]; then
    echo java -classpath ${BYTEMAN_SUBMIT_JAR}:${BYTEMAN_JAR}:${HELPERSDT_JAR} org.jboss.byteman.agent.submit.Submit -p $bmport $bmcmd $btmfile
fi

exec java -classpath ${BYTEMAN_SUBMIT_JAR}:${BYTEMAN_JAR}:${HELPERSDT_JAR} org.jboss.byteman.agent.submit.Submit -p $bmport $bmcmd $btmfile
