#!/bin/bash
function test_mgr() {
    whost=$1
    if [[ "$whost" == *:* ]] && [[ "$whost" != *[* ]] ; then
        whost="[$whost]"
    fi
    if curl -gs https://${whost}/confluent-api/ > /dev/null; then
        return 0
    fi
    return 1
}

function initconfluentscriptstmp() {
    if [ -z "$confluentscripttmpdir" ]; then
        mkdir -p /opt/confluent/tmpexec
        confluentscripttmpdir=$(mktemp -d /opt/confluent/tmpexec/confluentscripts.XXXXXXXXX)
    fi
}

function confluentpython() {
    if [ -x /usr/libexec/platform-python ]; then
        /usr/libexec/platform-python $*
    elif [ -x /usr/bin/python3 ]; then
        /usr/bin/python3 $*
    elif [ -x /usr/bin/python ]; then
        /usr/bin/python $*
    elif [ -x /usr/bin/python2 ]; then
        /usr/bin/python2 $*
    fi
}

function set_confluent_vars() {
    if [ -z "$nodename" ]; then
        nodename=$(grep ^NODENAME: /etc/confluent/confluent.info | awk '{print $2}')
    fi
    if [[ "$confluent_mgr" == *"%"* ]]; then
        confluent_mgr=""
    fi
    if [ -z "$confluent_mgr" ]; then
        confluent_mgr=$(grep ^deploy_server: /etc/confluent/confluent.deploycfg | sed -e 's/[^ ]*: //')
        if ! test_mgr $confluent_mgr; then
            confluent_mgr=$(grep ^deploy_server_v6: /etc/confluent/confluent.deploycfg | sed -e 's/[^ ]*: //')
            if [[ "$confluent_mgr" = *":"* ]]; then
                confluent_mgr="[$confluent_mgr]"
            fi
        fi
        if ! test_mgr $confluent_mgr; then
            BESTMGRS=$(grep ^EXTMGRINFO: /etc/confluent/confluent.info | grep '|1$' | sed -e 's/EXTMGRINFO: //' -e 's/|.*//')
            OKMGRS=$(grep ^EXTMGRINFO: /etc/confluent/confluent.info | grep '|0$' | sed -e 's/EXTMGRINFO: //' -e 's/|.*//')
            for confluent_mgr in $BESTMGRS $OKMGRS; do
                if [[ $confluent_mgr == *":"* ]]; then
                    confluent_mgr="[$confluent_mgr]"
                fi
                if test_mgr $confluent_mgr; then
                    break
                fi
            done
        fi
    fi
    if [ -z "$confluent_profile" ]; then
        confluent_profile=$(grep ^profile: /etc/confluent/confluent.deploycfg | sed -e 's/[^ ]*: //')
    fi
    export confluent_profile confluent_mgr nodename
}

fetch_remote() {
    curlargs=""
    if [ -f /etc/confluent/ca.pem ]; then
        curlargs=" --cacert /etc/confluent/ca.pem"
    fi
    set_confluent_vars
    mkdir -p $(dirname $1)
    whost=$confluent_mgr
    if [[ "$whost" == *:* ]] && [[ "$whost" != *[* ]] ; then
        whost="[$whost]"
    fi
    curl -gf -sS $curlargs https://$whost/confluent-public/os/$confluent_profile/scripts/$1 > $1
    if [ $? != 0 ]; then echo $1 failed to download; return 1; fi
}

source_remote_parts() {
    unset confluentscripttmpdir
    initconfluentscriptstmp
    apiclient=/opt/confluent/bin/apiclient
    if [ -f /etc/confluent/apiclient ]; then
        apiclient=/etc/confluent/apiclient
    fi
    scriptlist=$(confluentpython $apiclient /confluent-api/self/scriptlist/$1|sed -e 's/^- //')
    for script in $scriptlist; do
        source_remote $1/$script
    done
    rm -rf $confluentscripttmpdir
    unset confluentscripttmpdir
}

run_remote_parts() {
    unset confluentscripttmpdir
    initconfluentscriptstmp
    apiclient=/opt/confluent/bin/apiclient
    if [ -f /etc/confluent/apiclient ]; then
        apiclient=/etc/confluent/apiclient
    fi
    scriptlist=$(confluentpython $apiclient /confluent-api/self/scriptlist/$1|sed -e 's/^- //')
    for script in $scriptlist; do
        run_remote $1/$script
    done
    rm -rf $confluentscripttmpdir
    unset confluentscripttmpdir
}

set_selinux_context() {
    if [ -x /usr/bin/chcon ]; then
        /usr/bin/chcon "$1" "$2"
    fi
}

source_remote() {
    set_confluent_vars
    unsettmpdir=0
    echo
    echo '---------------------------------------------------------------------------'
    echo Sourcing $1 from https://$confluent_mgr/confluent-public/os/$confluent_profile/scripts/
    initconfluentscriptstmp
    echo Sourcing from $confluentscripttmpdir
    cd $confluentscripttmpdir
    fetch_remote $1
    if [ $? != 0 ]; then echo $1 failed to download; return 1; fi
    chmod +x $1
    cmd=$1
    shift
    source ./$cmd
    cd - > /dev/null
    if [ "$unsettmpdir" = 1 ]; then
        rm -rf $confluentscripttmpdir
        unset confluentscripttmpdir
        unsettmpdir=0
    fi
    rm -rf $confluentscripttmpdir
    return $retcode
}

run_remote() {
    requestedcmd="'$*'"
    unsettmpdir=0
    set_confluent_vars
    echo
    echo '---------------------------------------------------------------------------'
    echo Running $requestedcmd from https://$confluent_mgr/confluent-public/os/$confluent_profile/scripts/
    if [ -z "$confluentscripttmpdir" ]; then
        unsettmpdir=1
    fi
    initconfluentscriptstmp
    echo Executing in $confluentscripttmpdir
    cd $confluentscripttmpdir
    fetch_remote $1
    if [ $? != 0 ]; then echo $requestedcmd failed to download; return 1; fi
    chmod +x $1
    cmd=$1
    set_selinux_context system_u:object_r:bin_t:s0 "$cmd"
    shift
    ./$cmd $*
    retcode=$?
    if [ $retcode -ne 0 ]; then
        echo "$requestedcmd exited with code $retcode"
    fi
    cd - > /dev/null
    if [ "$unsettmpdir" = 1 ]; then
        rm -rf $confluentscripttmpdir
        unset confluentscripttmpdir
        unsettmpdir=0
    fi
    return $retcode
}

run_remote_python() {
    echo
    set_confluent_vars
    if [ -f /etc/confluent/ca.pem ]; then
        curlargs=" --cacert /etc/confluent/ca.pem"
    fi
    echo '---------------------------------------------------------------------------'
    echo Running python script "'$*'" from https://$confluent_mgr/confluent-public/os/$confluent_profile/scripts/
    unset confluentscripttmpdir
    initconfluentscriptstmp
    echo Executing in $confluentscripttmpdir
    cd $confluentscripttmpdir
    mkdir -p $(dirname $1)
    whost=$confluent_mgr
    if [[ "$whost" == *:* ]] && [[ "$whost" != *[* ]] ; then
        whost="[$whost]"
    fi
    curl -gf -sS $curlargs https://$whost/confluent-public/os/$confluent_profile/scripts/$1 > $1
    if [ $? != 0 ]; then echo "'$*'" failed to download; return 1; fi
    confluentpython $*
    retcode=$?
    echo "'$*' exited with code $retcode"
    cd - > /dev/null
    rm -rf $confluentscripttmpdir
    unset confluentscripttmpdir
    return $retcode
}

run_remote_config() {
    echo
    set_confluent_vars
    apiclient=/opt/confluent/bin/apiclient
    if [ -f /etc/confluent/apiclient ]; then
        apiclient=/etc/confluent/apiclient
    fi
    echo '---------------------------------------------------------------------------'
    echo Requesting to run remote configuration for "'$*'" from $confluent_mgr under profile $confluent_profile
    confluentpython $apiclient /confluent-api/self/remoteconfig/"$*" -d {}
    confluentpython $apiclient /confluent-api/self/remoteconfig/status -w 204
    echo
    echo 'Completed remote configuration'
    echo '---------------------------------------------------------------------------'
    return
}
startlegacyprogress() {
    [ -n "$progresspid" ] && return
    echo -en "Decrypting and extracting root filesystem: 0%\r"
    srcsz=$(du -sk /mnt/remote | awk '{print $1}')
    while [ -f /mnt/remoteimg/rootimg.sfs ]; do
        dstsz=$(du -sk /sysroot | awk '{print $1}')
        pct=$((dstsz * 100 / srcsz))
        if [ $pct -gt 99 ]; then
            pct=99
        fi
        echo -en "Decrypting and extracting root filesystem: $pct%\r"
        sleep 0.25
    done &
    progresspid=$!
}

extractmultisquash() {
    multisquashrc=1
    while read -r _ partsrc partmount; do
        normalizedmount=${partmount%/}
        [ "$normalizedmount" = /mnt/remote ] || continue
        if unsquashfs -force -d /sysroot "$partsrc"; then
            multisquashrc=0
        fi
        break
    done < /tmp/mountparts.sh
    [ $multisquashrc -eq 0 ] || return 1
    while read -r _ partsrc partmount; do
        normalizedmount=${partmount%/}
        [ "$normalizedmount" = /mnt/remote ] && continue
        partdest=/sysroot${normalizedmount#/mnt/remote}
        mkdir -p "$partdest"
        if ! unsquashfs -force -d "$partdest" "$partsrc"; then
            return 1
        fi
    done < /tmp/mountparts.sh
}

cleanupremotemounts() {
    if [ -f /tmp/mountparts.sh ]; then
        for partmount in $(awk '{ mounts[NR] = $3 } END { for (idx = NR; idx > 0; idx--) print mounts[idx] }' /tmp/mountparts.sh); do
            umount "$partmount"
        done
        for partdev in $(awk '{ devices[NR] = $NF } END { for (idx = NR; idx > 0; idx--) print devices[idx] }' /tmp/setupmount.sh); do
            dmsetup remove "$partdev"
        done
    else
        umount /mnt/remote
    fi
}

# Extract an untethered/uncompressed diskless root image into /sysroot and
# tear down the source devices; expects rootimgformat, mountsrc, and loopdev
# from imageboot.sh and the image content mounted under /mnt/remote.
extract_untethered_rootimg() {
    extractrc=1
    progresspid=
    if [ "$rootimgformat" = squashfs ] && command -v unsquashfs > /dev/null 2>&1; then
        echo "Decrypting and extracting root filesystem in parallel"
        if unsquashfs -force -d /sysroot "$mountsrc"; then
            extractrc=0
        else
            echo "Parallel root filesystem extraction failed, retrying with cp"
        fi
    elif [ "$rootimgformat" = confluent_multisquash ] && command -v unsquashfs > /dev/null 2>&1; then
        echo "Decrypting and extracting multipart root filesystem in parallel"
        if extractmultisquash; then
            extractrc=0
        else
            echo "Parallel multipart root filesystem extraction failed, retrying with cp"
        fi
    fi
    if [ $extractrc -ne 0 ]; then
        rm -rf /sysroot/* /sysroot/.[!.]* /sysroot/..?*
        startlegacyprogress
        if cp -a /mnt/remote/. /sysroot/; then
            extractrc=0
        fi
    fi
    if [ $extractrc -ne 0 ]; then
        if [ -n "$progresspid" ]; then
            kill "$progresspid" 2> /dev/null || true
            wait "$progresspid" 2> /dev/null || true
            progresspid=
        fi
        echo "Failed to extract the root filesystem"
    fi
    cleanupremotemounts
    if [ -e /dev/mapper/cryptimg ]; then
        dmsetup remove cryptimg
    fi
    losetup -d $loopdev
    rm /mnt/remoteimg/rootimg.sfs
    umount /mnt/remoteimg
    if [ -n "$progresspid" ]; then
        wait "$progresspid"
        echo -e "Decrypting and extracting root filesystem: 100%"
    fi
    return $extractrc
}

#If invoked as a command, use the arguments to actually run a function
(return 0 2>/dev/null) || $1 "${@:2}"