#!/usr/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"
    set -x
else
    exec >/dev/null
    # NB: preserve stderr
fi

prefix=/usr
exec_prefix=/usr
pkglibexecdir=/usr/libexec/systemtap

HELPERSDT_JAR=${pkglibexecdir}/HelperSDT.jar
if [ ! -f ${HELPERSDT_JAR} ]; then
    exec 1>&2
    echo "Missing $HELPERSDT_JAR"
    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

    bminstall -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
    bmsubmit -p $bmport -s $HELPERSDT_JAR
    if [ $? -ne 0 ]; then
        echo "Byteman agent failed to load HelperSDT.jar 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

bmsubmit -p $bmport $bmcmd $btmfile
