#!/bin/bash
set -eu

COMMAND=($@)
QGA_SOCKET=/var/run/netvm_qga.sock
LOCKFILE=/var/run/netvm.pid

json() {
        jq -ncM "$@"
}

qemu_execute() {
        local COMMAND ARGS
        COMMAND="$1"
        ARGS="${2-}"

        json --arg cmd "$COMMAND" --argjson args "$ARGS" '{"execute": $cmd, "arguments": $args}' >&$FD_SOCKET_OUT

        local LINE
        read -t 5 -r -u $FD_SOCKET_IN LINE

        local ERROR=$(jq -r '.error.desc // empty' <<< "$LINE")
        if [[ -n "$ERROR" ]]; then
                echo "$ERROR" >&2
                return 1
        fi

        GA_RETURN=$(jq -cM .return <<< "$LINE")
}

qemu_ga() {
        local COMMAND ARGS
        COMMAND="$1"
        ARGS="$2"

        coproc FDS (
                exec socat - UNIX-CONNECT:"${QGA_SOCKET}"
        )

        FD_SOCKET_IN=${FDS[0]}
        FD_SOCKET_OUT=${FDS[1]}

        local PID=$$
        qemu_execute guest-sync "$(json --argjson pid "$PID" '{"id": $pid}')"
        [[ "$(jq -re . <<< "$GA_RETURN")" = "$$" ]] || (echo "guest-sync mismatch" >&2 && return 1)

        qemu_execute "$COMMAND" "$ARGS"
        echo "$GA_RETURN" 2>&1

        local RETURN
        kill -INT "$FDS_PID" 2>/dev/null
        wait "$FDS_PID" || RETURN=$?
        if [[ $RETURN != 130 ]]; then
                return $RETURN
        fi
}

function cmd_start(){
        [ ! -f "${LOCKFILE}" ] || { echo "Error: Netvm already running"; exit 1; }
        echo "Starting netvm";
        qemu-system-x86_64 -m 512M \
                --machine q35 \
                -nographic \
                -serial none \
                -monitor none \
                -net none \
                -cdrom /guest.img \
                -boot order=d \
                -chardev socket,path=/var/run/netvm_qga.sock,server=on,wait=off,id=qga0 \
                -device virtio-serial \
                -device virtserialport,chardev=qga0,name=org.qemu.guest_agent.0 &
        echo $! > "${LOCKFILE}"
}

function cmd_stop(){
        pkill -F "${LOCKFILE}" && rm "${LOCKFILE}"
}

function cmd_status(){
        qemu_ga guest-get-host-name "{}" | jq -r '."host-name"'
        pid=$(qemu_ga guest-exec '{"path": "uptime", "capture-output": true}' | jq -r '.pid')
        out=$(qemu_ga guest-exec-status "$(jq -n --argjson pid "$pid" '{pid: $pid }')" \
                | jq -r '."out-data"' \
                | base64 -d \
        )
        echo $out
}

function cmd_push(){
        local source="${COMMAND[1]}"
        local dest="${COMMAND[2]}"
        fo_request=$(jq -n --arg dest "$dest" '{"path": $dest, "mode": "w" }')
        handle=$(qemu_ga guest-file-open "$fo_request")
        bufb64=$(base64 "$source")
        count=$(cat "$source" | wc -c)
        fw_request=$(jq -n \
                --argjson handle $handle \
                --argjson count $count \
                --arg bufb64 "$bufb64" \
                '{handle: $handle, "buf-b64": $bufb64, count: $count }' \
        )
        qemu_ga guest-file-write "$fw_request"
        fh_request=$(jq -n --argjson handle $handle '{handle: $handle}' )
        qemu_ga guest-file-flush "$fh_request"
        qemu_ga guest-file-close "$fh_request"
}

function cmd_pull(){
        local source="${COMMAND[1]}"
        local dest="${COMMAND[2]}"
        fo_request=$(jq -n --arg source "$source" '{"path": $source}')
        handle=$(qemu_ga guest-file-open "$fo_request")
        fr_request=$(jq -n \
                --argjson handle $handle \
                '{handle: $handle, count: 48000000 }' \
        )
        out=$(qemu_ga guest-file-read "$fr_request")
        echo $out | jq -r '."buf-b64"' | base64 -d > $dest
}

function cmd_run(){
        [ -z "${COMMAND[1]}" ] && { echo "Error: missing command"; exit 1; }
        [ -f "${LOCKFILE}" ] || { echo "Error: Netvm is not running"; exit 1; }
        [ -S "${QGA_SOCKET}" ] || { echo "Error: Netvm QGA socket is missing"; exit 1; }
        local cmd="${COMMAND[1]}"
        local args="${COMMAND[@]:2}"
        local args_json="[]"
        if [[ -n "$args" ]]; then
                args_json=$(printf '%s\n' "$args" | jq -R . | jq -s .)
        fi
        local request
        request=$( \
                jq -n \
                        --arg path "$cmd" \
                        --argjson args "$args_json" \
                        '{
                                        path: $path,
                                        arg: $args,
                                        "capture-output": true
                        }' \
        )
        pid=$(qemu_ga guest-exec "$request" | jq -r '.pid')
        local exited=false
        until [ "$exited" == "true" ]; do \
                out=$(qemu_ga guest-exec-status "$(jq -n --argjson pid "$pid" '{pid: $pid }')" )
                exited=$(echo $out | jq -r '.exited')
                if $exited && jq -r 'has("out-data")' >/dev/null < <(echo $out); then
                        echo "$out" | jq -r '."out-data"' | base64 -d
                        break
                fi
                sleep 1
        done
}

cmd_usage() {
	cat <<-_EOF
	netvm
	
	Control network vm headlessly via QMP protocol

	Usage:
	  netvm start
	      Start headless network vm in the background
	  netvm stop
	      Stop headless network vm
	  netvm status
	      Get hostname and uptime from running network vm
	  netvm push <local-path> <remote-path>
	      Push a local file to the network VM
	  netvm pull <remote-path> <local-path>
	      Pull a file from the network VM
	  netvm run "<command>"
	      Run a command in network vm and get stdout
	_EOF
}

case "$1" in
        status) shift; cmd_status $@ ;;
        start)  shift; cmd_start  $@ ;;
        stop)   shift; cmd_stop   $@ ;;
        push)   shift; cmd_push   $@ ;;
        pull)   shift; cmd_pull   $@ ;;
        run)    shift; cmd_run    $@ ;;
        help)   shift; cmd_usage  $@ ;;
        *)             cmd_usage  $@ ;;
esac