#!/usr/bin/env bash
set -o nounset
# set -o xtrace

###############################################################################

##############
# Constants #
############

readonly NAME="tcpsnitch"
readonly VERSION=0.1
readonly VERSION_STR="${NAME} version ${VERSION}"
#readonly SERVER="192.168.99.100:3000"
readonly SERVER="https://tcpsnitch.org"
#readonly SERVER="staging.tcpsnitch.org"
readonly SERVER_ENDPOINT="${SERVER}/app_traces"
readonly I386_LIB="lib${NAME}.so.${VERSION}-i386"
readonly AMD64_LIB="lib${NAME}.so.${VERSION}-x86-64"
readonly ARM_LIB="lib${NAME}.so.${VERSION}-arm"
readonly SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
readonly PROP_PREFIX="be.ucl.${NAME}"
readonly HIDDEN_ERRORS="wrong ELF class"
readonly TOP_PID=$$
readonly CWD=$(pwd)

# Option default values
OPT_A=0
OPT_B=0
OPT_C=0
OPT_D=""
OPT_F=2
OPT_L=1
OPT_N=0
OPT_P=0
OPT_T=1000
OPT_U=0
OPT_V=0

# Options saved in meta files
META_OPTIONS_NAMES=(opt_b opt_f opt_u)
META_OPTIONS_COUNT=${#META_OPTIONS_NAMES[@]}
META_OPTIONS_VALS=($OPT_B $OPT_F $OPT_U)

###############################################################################

############
# Options #
##########

usage() {
    local _head="Usage: ${NAME}"
    local _skip=$(printf "%0.s " $(seq 1 ${#_head}))
    echo "${_head} [-achpv] [ -b <bytes> ] [ -d <dir>] [ -f <lvl> ]"
    echo "${_skip} [ -k <pkg> ] [ -l <lvl> ] [ -t <msec> ]"
    echo "${_skip} [ -u <usec> ] [ --version ] <app> [<args>]"
    echo ""
    echo "<app>       cmd/package to spy on."
    echo "<args>      args to <app>."
    echo "-a          instrument & launch app on connected android device."
    echo "-b <bytes>  dump tcp_info every <bytes> (0 means NO dump, def 0)."
    echo "-c          activate capture of pcap traces (only on Linux)."
    echo "-d <dir>    dir to save traces (defaults to random dir in /tmp)."
    echo "-f <lvl>    verbosity of logs to file (0 to 5, defaults to 2)."
    echo "-h          show this help text."
    echo "-k <pkg>    kill instrumented android <pkg> and pull traces."
    echo "-l <lvl>    verbosity of logs to stderr (0 to 5, defaults to 2)."
    echo "-n          do (n)ot send traces to web server."
    echo "-p          pedantic, ask a lot of annoying questions."
    echo "-t <msec>   dump to JSON file every <msec> (def. 1000)."
    echo "-u <usec>   dump tcp_info every <usec> (0 means NO dump, def 0)."
    echo "-v          activate verbose output (not really implemented)."
    echo "--version   print ${NAME} version."
}

parse_options() {
    # Parse options
    while getopts ":achnpvb:d:f:k:l:t:u:-:" opt; do
        case "${opt}" in
            -) # Trick to parse long options with getopts.
                case "${OPTARG}" in
                    version)
                        info "${VERSION_STR}"
                        exit 0
                        ;;
                esac
                ;;
            a)
                OPT_A=1
                ;;
            b)
                assert_int "${OPTARG}" "invalid -b argument: '${OPTARG}'" 
                OPT_B=${OPTARG}
                ;;
            c)
                OPT_C=1;
                ;;
            d)
                if [[ ! -d "${OPTARG}" ]] ; then
                    error "invalid -d argument: '${OPTARG}'"
                fi
                OPT_D=$(readlink -f "$OPTARG")
                ;;
            f)
                assert_int "${OPTARG}" "invalid -f argument: '${OPTARG}'" 
                OPT_F=${OPTARG}
                ;;
            h)
                usage
                exit 0
                ;;
            k)
                tcpsnitch_android_teardown $@
                exit 0
                ;;
            l)
                assert_int "${OPTARG}" "invalid -l argument: '${OPTARG}'" 
                OPT_L=${OPTARG}
                ;;
            n)
                OPT_N=1
                ;;
            p)
                OPT_P=1
                ;;
            u)
                assert_int "${OPTARG}" "invalid -u argument: '${OPTARG}'" 
                OPT_U=${OPTARG}
                ;;
            t)
                assert_int "${OPTARG}" "invalid -t argument: '${OPTARG}'"
                OPT_T=${OPTARG}
                ;;
            v)
                OPT_V=$((OPT_V+1))
                ;;
            \?)
                error "invalid option"
                ;;
        esac
    done
}

###############################################################################

############
# Helpers #
##########

error() {
    declare msg="$1"
    echo "${NAME}: ${msg}."
    echo "Try '${NAME} -h' for more information."
    kill -s TERM $TOP_PID
}

info() {
    echo "${1}."
}

is_integer() {
    [[ "$1" =~ ^[0-9]+$ ]]
}

assert_int() {
    declare int="$1"
    declare error_msg="$2"
    if ! is_integer "$int"; then
        error "$error_msg"
    fi
}

cd_script_dir() {
    cd "$SCRIPT_DIR" || exit "Could not cd to ${SCRIPT_DIR}"
}

assert_lib_present() {
    declare lib="$1"
    if [[ ! -f "$lib" ]]; then
        error "${lib} missing! Reinstall ${NAME} (see README)"
    fi
}

validate_args_number() {
    if [[ $# -lt 1 ]]; then
        error "too few arguments"
    fi
}

create_opt_d_dir() {
    if [[ -z "$OPT_D" ]]; then
        OPT_D=$(mktemp -d)
        chmod 777 "$OPT_D"
    fi
}

zip_trace() {
    cd "${OPT_D}" || error "Could not cd to ${OPT_D}"
    tar -czf archive.tar.gz ./*
}

upload_trace() {
    if [[ $OPT_N -eq "1" ]]; then exit; fi

    # Test if trace is empty
    if ! ls ${OPT_D}/*/*.json >/dev/null 2>/dev/null; then
        error "Nothing to trace. Please report a bug if you have reasons to believe the trace should not be empty (https://github.com/GregoryVds/tcpsnitch/issues)"
    fi

    zip_trace

    declare form="app_trace[archive]=@${OPT_D}/archive.tar.gz"

    if [[ $OPT_P -eq "1" ]]; then
        read -p "Workload decription (e.g. 'Opened app & navigated to google.com'): " workload

        PS3='Select connectivity type: '
        options=("wifi" "ethernet" "lte")
        select connectivity in "${options[@]}"; do
            case "$connectivity" in
                "wifi"|"ethernet"|"lte")
                    break
                    ;;
                *)
                    echo "Invalid option"
            esac
        done

        info "Uploading trace..."
        if ! curl --connect-timeout 2 -F "$form" -F "app_trace[workload]=$workload" -F "app_trace[connectivity]=$connectivity" "$SERVER_ENDPOINT"; then
            info "Failed to connect to ${SERVER}"
        fi
    else
        info "Uploading trace..."
        if ! curl --connect-timeout 2 -F "$form" "$SERVER_ENDPOINT"; then
            info "Failed to connect to ${SERVER}"
        fi
    fi
}

###############################################################################

####################
# Android helpers #
##################

get_android_package() {
    declare package_name="$1"

    if [[ $# -ne 1 ]]; then
        error "missing package argument"
    fi

    # Check that adb sees exactly 1 device
    if [ $(adb devices | wc -l) -ne 3 ]; then
        error "adb must see exactly 1 device"
    fi

    declare prefix="package:"
    declare match=$(adb shell pm list packages | grep -m 1 -s "${package_name}" | tr -d '\r')

    # Check that package exists
    if [ -z "${match}" ]; then
        error "invalid argument: package '${package_name}' not found"
    fi

    PACKAGE=${match#${prefix}}
    info "Found Android package: '${PACKAGE}'"
}

start_android_package() {
    adb shell "monkey -p ${PACKAGE} -c android.intent.category.LAUNCHER 1 >/dev/null"
    info "Start package '${PACKAGE}'"
}

kill_android_package() {
    adb shell su -c am force-stop "${PACKAGE}"
}

pull_trace_for_android_package() {
    declare tracesdir="/data/data/${PACKAGE}/${NAME}"
    exists=$(adb shell "test -d ${tracesdir}; echo \$?" | tr -d '\r')
    if [[ "$exists" -eq "1" ]] ; then
        error "no traces for '${PACKAGE}'"
    fi

    info "Pulling trace from Android device..."
    adb shell su -c chown -R shell "$tracesdir"
    adb pull "$tracesdir" "$OPT_D" 2>/dev/null
    adb shell su -c rm -rf "$tracesdir"
    info "Trace saved in $OPT_D"
}

###############################################################################

##########
# Linux #
########

tcpsnitch_linux() {
    if [ -f "$cmd" ]; then
        # Resolve absolute path for $1 & update positional parameter accordingly
        cmd=$(readlink -f "$cmd")
        set -- "$cmd" "${@:2}"

        if ! [[ -x "$cmd" ]]; then
            error "invalid cmd: '$cmd' is not executable"
        fi
    else
        if ! which "$cmd" > /dev/null; then
            error "invalid argument: '$cmd' is not in \$PATH"
        fi
    fi

    # Test if libs are present
    cd_script_dir
    readonly ENABLE_I386=$(cat enable_i386 2>/dev/null || false)
    $ENABLE_I386 && assert_lib_present "$I386_LIB"
    assert_lib_present "$AMD64_LIB"
    readonly LINUX_GIT_HASH=$(cat linux_git_hash)

    create_opt_d_dir

    # Extract app (split on '/' and take last)
    # Global replace (//) in $1 of / (escaped) by ' ' & interpret as array
    declare array=(${1//\// });
    declare app=${array[-1]}

    # Write meta data
    declare meta_dir="$OPT_D/meta"
    mkdir "$meta_dir"
    echo "$app" > "${meta_dir}/app"
    echo "$@" > "${meta_dir}/cmd"
    uname -r > "${meta_dir}/kernel"
    uname -m > "${meta_dir}/machine"
    sysctl net > "${meta_dir}/net" 2>/dev/null
    uname -s > "${meta_dir}/os"
    echo "$VERSION" > "${meta_dir}/version"
    echo "$LINUX_GIT_HASH" > "${meta_dir}/git_hash"
    if ! ip link | grep 'link/ether' | awk '{print $2}' | sha256sum | awk '{print $1}' > "${meta_dir}/host_id"; then
        ifconfig | grep 'HWaddr' | awk '{print $NF}' | sha256sum | awk '{print $1}' > "${meta_dir}/host_id"
    fi
    "$1" --version 2>/dev/null >"${meta_dir}/app_version"
    for i in $(seq 0 $((${META_OPTIONS_COUNT}-1))); do
        echo "${META_OPTIONS_VALS[$i]}" > ${meta_dir}/${META_OPTIONS_NAMES[$i]}
    done

    local _preload_opt=""
    ${ENABLE_I386} && _preload_opt+="$(readlink -f "$I386_LIB") "
    _preload_opt+=$(readlink -f "$AMD64_LIB")

    cd "$CWD"
    # libtcpsnitch uses fd 3 & 4 as stdout & stderr repectively. This allows to
    # distinguish the standard outputs of tcpsnitch and of the traced process.
    # The distinction is however not used at the moment s standard streams of
    # libtcpsnitch & the traced process are currently sent to the same place.
    exec 3>&1
    exec 4>&2

    {\
    TCPSNITCH_OPT_B=$OPT_B \
    TCPSNITCH_OPT_C=$OPT_C \
    TCPSNITCH_OPT_D=$OPT_D \
    TCPSNITCH_OPT_F=$OPT_F \
    TCPSNITCH_OPT_L=$OPT_L \
    TCPSNITCH_OPT_T=$OPT_T \
    TCPSNITCH_OPT_U=$OPT_U \
    TCPSNITCH_OPT_V=$OPT_V \
    LD_PRELOAD="${_preload_opt}" "$@" 1>&3; \
    # Filter out some errors
    } 2>&1 | grep -E -v "$HIDDEN_ERRORS" 1>&2

    info "Trace saved in ${OPT_D}"

    upload_trace
}

###############################################################################

############
# Android #
##########

tcpsnitch_android() {
    cd_script_dir
    assert_lib_present "$ARM_LIB"
    readonly ANDROID_GIT_HASH=$(cat android_git_hash)
    get_android_package "$cmd"
    kill_android_package
    create_opt_d_dir

    # Desactivate Selinux
    adb shell su -c setenforce 0

    # Create logs dir
    declare LOGS_DIR="/data/data/${PACKAGE}/${NAME}"
    adb shell su -c rm -rf "$LOGS_DIR"
    adb shell su -c mkdir -m 777 -p "$LOGS_DIR"

    # Write meta data
    declare meta_dir="$LOGS_DIR/meta"
    adb shell mkdir "$meta_dir"
    adb shell "echo ${PACKAGE} > ${meta_dir}/app"
    adb shell "echo ${PACKAGE} > ${meta_dir}/cmd"
    adb shell "uname -r > ${meta_dir}/kernel"
    adb shell "uname -m > ${meta_dir}/machine"
    adb shell "sysctl net > ${meta_dir}/net 2>/dev/null"
    adb shell "echo android > ${meta_dir}/os"
    adb shell "echo ${VERSION} > ${meta_dir}/version"
    adb shell "echo ${PACKAGE} > ${meta_dir}/app"
    adb shell "echo ${ANDROID_GIT_HASH} > ${meta_dir}/git_hash"
    adb shell "ip link | grep 'link/ether' | awk '{print \$2}' | sha256sum | awk '{print \$1}' > ${meta_dir}/host_id"
    adb shell "dumpsys package ${PACKAGE} | grep versionName > ${meta_dir}/app_version"
    for i in $(seq 0 $((${META_OPTIONS_COUNT}-1))); do
        adb shell "echo ${META_OPTIONS_VALS[$i]} > ${meta_dir}/${META_OPTIONS_NAMES[$i]}"
    done

    # Upload lib
    LIBPATH="/data"
    adb shell su -c chmod 777 /data
    adb shell rm -f "${LIBPATH}/${ARM_LIB}"
    info "Uploading ${NAME} library to ${LIBPATH}/${ARM_LIB}"
    adb push ${ARM_LIB} ${LIBPATH} 2>/dev/null

    # Properties are limited to 32 chars includind the NULL byte.
    # With "wrap." being 5 chars, we have 26 chars left the app name.
    adb shell setprop wrap."${PACKAGE:0:26}" LD_PRELOAD="${LIBPATH}/${ARM_LIB}"
    adb shell setprop "${PROP_PREFIX}.opt_b" "$OPT_B"
    adb shell setprop "${PROP_PREFIX}.opt_d" "$LOGS_DIR"
    adb shell setprop "${PROP_PREFIX}.opt_f" "$OPT_F"
    adb shell setprop "${PROP_PREFIX}.opt_l" "$OPT_L"
    adb shell setprop "${PROP_PREFIX}.opt_t" "$OPT_T"
    adb shell setprop "${PROP_PREFIX}.opt_u" "$OPT_U"
    adb shell setprop "${PROP_PREFIX}.opt_v" "$OPT_V"

    # Those properties are used by this bash script only. We set them to
    # retrieve them on -k.
    adb shell setprop "${PROP_PREFIX}.opt_d_k" "$OPT_D"
    adb shell setprop "${PROP_PREFIX}.opt_n_k" "$OPT_N"
    adb shell setprop "${PROP_PREFIX}.opt_p_k" "$OPT_P"

    start_android_package
    info "Execute './${NAME} -k ${cmd}' to terminate the capture"
}

tcpsnitch_android_teardown() {
    OPT_D=$(adb shell getprop "${PROP_PREFIX}.opt_d_k" | tr -d '\r')
    OPT_N=$(adb shell getprop "${PROP_PREFIX}.opt_n_k" | tr -d '\r')
    OPT_P=$(adb shell getprop "${PROP_PREFIX}.opt_p_k" | tr -d '\r')

    get_android_package "${2}"
    kill_android_package
    pull_trace_for_android_package
    upload_trace
}

###############################################################################

################
# MAIN SCRIPT #
##############

trap "exit 1" TERM

parse_options "$@"
shift $((OPTIND - 1))

validate_args_number "$@"
cmd="$1"

if [[ $OPT_A -eq "1" ]]; then
    tcpsnitch_android "$@"
else
    tcpsnitch_linux "$@"
fi
