#!/bin/bash

# muxsa-kvm2png 
#
# part of muxsa, https://git-nks-public.tik.uni-stuttgart.de/edu/muxsa
#
# this script makes a series of screenshots from a KVM virtual machine
# and saves them as PNG files.  It will press the PgDown key after each
# screenshot.  Intended usage: dump a presentation (e.g., Powerpoint 
# running in a Windows VM) into a series of PNG files, for further 
# processing with the muxsa tool chain.

# MIT License
# 
# Copyright (c) 2023 Sebastian Kiesel <sebastian.kiesel@tik.uni-stuttgart.de>
# 
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.

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

echo "This is muxsa-kvm2png."

# various parameters can be set as an environment variable (export XX="y z")
# or can be set in our RC file. Otherwise we will assign default values below.

if [ -r "${HOME}/.muxsarc" ] ; then
    echo "Found ${HOME}/.muxsarc ... will source it."
    . "${HOME}/.muxsarc"
fi

echo "Will work with these settings:"

echo "MUXSA_KVM2PNG_VM_NAME=${MUXSA_KVM2PNG_VM_NAME:="win10-office"}"
echo "MUXSA_KVM2PNG_SIZE_X=${MUXSA_KVM2PNG_SIZE_X:="1920"}"
echo "MUXSA_KVM2PNG_SIZE_Y=${MUXSA_KVM2PNG_SIZE_Y:="1080"}"
echo "MUXSA_KVM2PNG_EXTRA_X=${MUXSA_KVM2PNG_EXTRA_X:="0"}"
echo "MUXSA_KVM2PNG_EXTRA_Y=${MUXSA_KVM2PNG_EXTRA_Y:="0"}"
echo "MUXSA_KVM2PNG_SLIDE_PREFIX=${MUXSA_KVM2PNG_SLIDE_PREFIX:=""}"
echo "MUXSA_KVM2PNG_SLIDE_NUMBER_START=${MUXSA_KVM2PNG_SLIDE_NUMBER_START:="0"}"
echo "MUXSA_KVM2PNG_SLIDE_NUMBER_END=${MUXSA_KVM2PNG_SLIDE_NUMBER_END:="9999"}"
echo "MUXSA_KVM2PNG_FORCE_OVERWRITE=${MUXSA_KVM2PNG_FORCE_OVERWRITE:=""}"
echo "MUXSA_KVM2PNG_BUGFIX_SHOOT_TWICE=${MUXSA_KVM2PNG_BUGFIX_SHOOT_TWICE:=""}"
echo "MUXSA_KVM2PNG_SLEEP=${MUXSA_KVM2PNG_SLEEP:="3"}"
echo "MUXSA_KVM2PNG_MD5_LAST_SLIDE=${MUXSA_KVM2PNG_MD5_LAST_SLIDE:="092eb68a91b4d2a73833f00a01dc5cb0"}"
echo

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

TMP_FILE="$(mktemp -t muxsa-kvm2png-XXXXXXXXXX.ppm)"
if [ ! -w "${TMP_FILE}" ] ; then
    echo "Error: cannot create tmp_file. abort."
    exit 1
fi
trap 'rm -f -- "${TMP_FILE}"' EXIT

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

muxsa-kvm2png-screenshot-to-tmp () {
    echo "muxsa-kvm2png: screenshot ${MUXSA_KVM2PNG_VM_NAME} ${TMP_FILE}"
    virsh -c qemu:///system "screenshot ${MUXSA_KVM2PNG_VM_NAME} ${TMP_FILE}"
    if [ $? -ne 0 ] ; then
        echo "muxsa-kvm2png: screenshot failed. abort."
        exit 1
    fi

    # the first sceenshot after a resolution change is sometimes garbled
    # workaround: save it twice
    if [ "${MUXSA_KVM2PNG_BUGFIX_SHOOT_TWICE}" ] ; then
        virsh -c qemu:///system \
            "screenshot ${MUXSA_KVM2PNG_VM_NAME} ${TMP_FILE}"
        if [ $? -ne 0 ] ; then
            echo "muxsa-kvm2png: screenshot failed. abort."
            exit 1
        fi
    fi

    # get the size in pixels of the screenshot we have taken
    # (needs identify from imagemagick), and ...
    TMP_SZ_X=$(identify -format "%w" "${TMP_FILE}")
    TMP_SZ_Y=$(identify -format "%h" "${TMP_FILE}")

    # ... compare to desired size
    if  [ ${TMP_SZ_X} -eq ${MUXSA_KVM2PNG_SIZE_X} -a \
          ${TMP_SZ_Y} -eq ${MUXSA_KVM2PNG_SIZE_Y} ] ; then
        echo "Screenshot is ${TMP_SZ_X} x ${TMP_SZ_Y} - as desired. good."
        TMP_SZ_OK=1
    else
        echo "Warning: Screenshot is ${TMP_SZ_X} x ${TMP_SZ_Y} but desired"\
            "size is ${MUXSA_KVM2PNG_SIZE_X} x ${MUXSA_KVM2PNG_SIZE_Y} ."
        TMP_SZ_OK=""
    fi
}

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

echo "Performing a test screenshot to check image size."
muxsa-kvm2png-screenshot-to-tmp
if [ ! "${TMP_SZ_OK}" ] ; then
    echo "We will try to adjust the size of the spicy window."
    case "$(read -p "Press q [return] to abort or [return] to continue: ")" in
        n|N|q|Q)
            echo "muxsa-kvm2png aborted."
            exit 1;
            ;;
    esac
    xdotool search --class spicy windowsize %@ \
        $(( ${MUXSA_KVM2PNG_SIZE_X} + ${MUXSA_KVM2PNG_EXTRA_X} )) \
        $(( ${MUXSA_KVM2PNG_SIZE_Y} + ${MUXSA_KVM2PNG_EXTRA_Y} ))
    sleep 5
    echo "Performing another test screenshot to check image size."
    muxsa-kvm2png-screenshot-to-tmp
    if [ ! "${TMP_SZ_OK}" ] ; then
        echo "Screenshot still does not have the desired size in pixels."
        echo "Maybe, resizing the spicy window has failed."
        echo "Or maybe, you need to adjust MUXSA_KVM2PNG_EXTRA_X / _Y"
        # extra size accounts for menus of the spicy window.
        # if the desired size is larger than the actual size, that
        # difference has to be added to the current extra size
        echo "Current MUXSA_KVM2PNG_EXTRA_X value: ${MUXSA_KVM2PNG_EXTRA_X}"
        echo "Try: $(( ${MUXSA_KVM2PNG_SIZE_X} - ${TMP_SZ_X} + \
                ${MUXSA_KVM2PNG_EXTRA_X} ))"
        echo "Current MUXSA_KVM2PNG_EXTRA_Y value: ${MUXSA_KVM2PNG_EXTRA_Y}"
        echo "Try: $(( ${MUXSA_KVM2PNG_SIZE_Y} - ${TMP_SZ_Y} + \
                ${MUXSA_KVM2PNG_EXTRA_Y} ))"
        echo "You can set these as environment variables or in ~/.muxsarc"
        echo "muxsa-kvm2png aborted."
        exit 1;
    fi
fi
echo "Successfully took test screenshot with the right size."

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

cat <<EOF

This is muxsa-kvm2png.

Name of the virtual machine: ${MUXSA_KVM2PNG_VM_NAME}
Screenshot size ${MUXSA_KVM2PNG_SIZE_X}x${MUXSA_KVM2PNG_SIZE_Y} verified.

Please press ... and return:

    q   to abort
    y   to start taking screenshots
    f   to send F5 to the virtual machine and start taking screenshots

EOF

case "$(read -t 10 -p "Answer: ")" in
    n|N|q|Q)
        echo "muxsa-kvm2png aborted."
        exit 1;
        ;;
    y|Y)
        ;;
    *)
        echo "Sending F5 to ${MUXSA_KVM2PNG_VM_NAME}"
        virsh -c qemu:///system "send-key ${MUXSA_KVM2PNG_VM_NAME} KEY_F5"
        sleep 3
        ;;
esac

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

SLIDE_NUMBER="${MUXSA_KVM2PNG_SLIDE_NUMBER_START:-0}"
while true ; do

    muxsa-kvm2png-screenshot-to-tmp
    if [ ! "${TMP_SZ_OK}" ] ; then
        echo "muxsa-kvm2png: screenshot failed. abort."
        exit 1
    fi

    OUTFILE=$(printf "${MUXSA_KVM2PNG_SLIDE_PREFIX}%04i.png" ${SLIDE_NUMBER})
    SLIDE_NUMBER=$(($SLIDE_NUMBER+1))
    if [ -f "${OUTFILE}" -a ! "${MUXSA_KVM2PNG_FORCE_OVERWRITE}" ] ; then
        echo "Outfile ${OUTFILE} exists and overwriting not allowed. abort."
        exit 1
    fi
    echo "Converting and saving screenshot to ${OUTFILE}"
    pnmtopng -force "${TMP_FILE}" > "$OUTFILE"
    if [ $? -ne 0 ] ; then
        echo "Converting and saving screenshot failed. abort."
        exit 1
    fi

    if [ "${MUXSA_KVM2PNG_MD5_LAST_SLIDE}" ] ; then
        if md5sum "${OUTFILE}" | \
            grep -q "^${MUXSA_KVM2PNG_MD5_LAST_SLIDE}" ; then

            echo "Screenshot's md5sum = MUXSA_KVM2PNG_MD5_LAST_SLIDE. end."
            break
        fi
    fi

    if [ "${SLIDE_NUMBER}" -gt "${MUXSA_KVM2PNG_SLIDE_NUMBER_END}" ] ; then
        echo "Maximum slide number reached. abort."
        exit 1
    fi

    echo "Sending PgDown and sleeping briefly"
    virsh -c qemu:///system "send-key ${MUXSA_KVM2PNG_VM_NAME} KEY_PAGEDOWN"
    sleep "${MUXSA_KVM2PNG_SLEEP}"

done
echo "muxsa-kvm2png finished successfully!"

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