#!/bin/bash

# enable debug output
[ -n "${L3VM_DEBUG}" ] && set -x

source /usr/lib/l3vm/l3vm_lib.sh

failed=()
L3VM_SSH_START_TIMEOUT=${L3VM_SSH_START_TIMEOUT:-120}
L3VM_POWEROFF_TIMEOUT=${L3VM_POWEROFF_TIMEOUT:-180}
L3VM_REPORT_EMAIL=${L3VM_REPORT_EMAIL:-l3vm-admin@localhost}
L3VM_LOG_FILE=${L3VM_LOG_FILE:-"/var/log/l3vmadm.log"}

allow_force_off=false
allow_running=false
stop_on_error=false
update=false
ssh_keys=false
change_password=false
error_report=false

L3VM_REPORT_TMP=$(mktemp /tmp/l3vmadm-report.XXXXXX.log)
trap 'rm -f "${L3VM_REPORT_TMP}"' EXIT

export L3VM_UPDATE_TIMEOUT=${L3VM_UPDATE_TIMEOUT:-60m}

# $1 host
# $2 remote operation
# SSH_EXTRA_ARGS
# sshpass_cmd
call_ssh() {
    sshpass_init

    timeout -k 30 ${L3VM_UPDATE_TIMEOUT} ${sshpass_cmd} ssh -o UserKnownHostsFile=/dev/null \
        -o StrictHostKeyChecking=no ${SSH_EXTRA_ARGS} root@"$1" "${SSH_CMD_TRACE:+set -x;} $2" </dev/null \
        | sed 's/^/  /'

    return $?
}

authorized_keys="$(cat ${L3VM_AUTHORIZED_KEYS:-/etc/l3vm/authorized_keys})"

authorize_snippet="
    echo ' == authorize snippet == '
    if [ -d /home/user/ ]; then
        mkdir -p /home/user/.ssh
        chown user:users /home/user/.ssh
        chmod 0700 /home/user/.ssh
        echo '$authorized_keys' | cat > /home/user/.ssh/authorized_keys
        chmod 0600 /home/user/.ssh/authorized_keys
        chown user:users /home/user/.ssh/authorized_keys
    fi
    mkdir -p /root/.ssh
    chmod 0700 /root/.ssh
    echo '$authorized_keys' | cat > /root/.ssh/authorized_keys
    chmod 0600 /root/.ssh/authorized_keys
    chown root:root /root/.ssh/authorized_keys
"

# define function in remote snippet
run_if_exists_def="
run_if_exists()
{

    local EXE=\${1}
    shift
    if [ -x \${EXE} ]; then
            echo \"* calling \${EXE} \$@\"
            eval \${EXE} \$@
    fi
}
"

btrfs_cleanup_snippet="
    echo ' == btrfs snippet == '
    $run_if_exists_def
    run_if_exists /usr/bin/snapper cleanup number
    run_if_exists /usr/share/btrfsmaintenance/btrfs-balance.sh '2>/dev/null | tail -n 8'
    run_if_exists /sbin/btrfs balance status -v /
"

# $1 vm name
# $2 host address
update_ssh_keys() {
    local RET=0
    inform "Copying authorized_keys to root & user on $1"
    call_ssh "${2}" "${authorize_snippet}"
    RET=$?
    if [ $RET -ne 0 ]; then
        warn "Updating ssh keys failed on $1 ($2) with return value $RET."
    fi
    inform "Done"
    return $RET
}

suseconnect_snippet="
    echo ' == SUSEConnect snippet == '
    $run_if_exists_def
    run_if_exists /usr/sbin/SUSEConnect --status
"
export suseconnect_snippet

update_9="
    online_update
"

update_10_11="
    zypper up -y --skip-interactive --auto-agree-with-licenses
"

update_12="
    CNT=6
    while [ \${CNT} -gt 0 ]; do
        ZYPP_LOCK_TIMEOUT=${L3VM_ZYPP_LOCK_TIMEOUT} zypper up -y --skip-interactive --auto-agree-with-licenses --download-as-needed
        RET=\$?
        CNT=\$((CNT-1))
        #  7 - ZYPPER_EXIT_ZYPP_LOCKED
        [ \${RET} -eq 7 ] || break
        sleep 30
    done
    $btrfs_cleanup_snippet
    $suseconnect_snippet
    exit \$RET"

update_micro="
    transactional-update dup --non-interactive --drop-if-no-change
    RET=\$?
    transactional-update cleanup
    $btrfs_cleanup_snippet
    $suseconnect_snippet
    exit \$RET"

update_tumbleweed="
    CNT=6
    while [ \${CNT} -gt 0 ]; do
        ZYPP_LOCK_TIMEOUT=${L3VM_ZYPP_LOCK_TIMEOUT} zypper dup -y --auto-agree-with-licenses --download-as-needed
        RET=\$?
        CNT=\$((CNT-1))
        #  7 - ZYPPER_EXIT_ZYPP_LOCKED
        [ \${RET} -eq 7 ] || break
        sleep 30
    done
    $btrfs_cleanup_snippet
    exit \$RET"

update_deb="
    CNT=6
    while [ \${CNT} -gt 0 ]; do
        apt-get -y update
        apt-get -y upgrade
        RET=\$?
        CNT=\$((CNT-1))
        [ \${RET} -eq 100 ] || break
        sleep 30
    done
    apt-get autoclean
    exit \$RET
"

update_res="
    yum -y update && yum -y upgrade
"

update_fedora="
    dnf upgrade --best -y
"

update_arch="
    pacman -Syu --noconfirm
    pacman -Sc --noconfirm
"

# select update function from VM name
# set update_fn variable to it
# $1 VM name
get_update_function()
{
    if [[ $1 =~ ^sle[sd]9 ]]; then
        inform "SLE9 - update_9() will be used"
        update_fn=${update_9}
    elif [[ $1 =~ ^sle[sd]1[01] ]] || \
            [[ $1 =~ ^opensuse-11 ]] || \
            [[ $1 =~ ^slepos11 ]] || \
            [[ $1 =~ ^slert10 ]] || \
            [[ $1 =~ ^slms1 ]]; then
        inform "SLE10+ based or openSUSE - update_10_11() will be used"
        update_fn=${update_10_11}
    elif [[ $1 =~ ^opensuse-tumbleweed ]]; then
        inform "OpenSUSE Tumbleweed - update_tumbleweed() will be used"
        update_fn=${update_tumbleweed}
    elif [[ $1 =~ ^sle[sd]1[256] ]] || \
            [[ $1 =~ ^suma ]] || \
            [[ $1 =~ ^opensuse-12 ]] || \
            [[ $1 =~ ^opensuse-leap ]]; then
        inform "SLE12 or newer - update_12() will be used"
        update_fn=${update_12}
    elif [[ $1 =~ ^sle-micro ]] || \
            [[ $1 =~ ^opensuse-microos ]] || \
            [[ $1 =~ ^alp ]]; then
        inform "SLE Micro - update_micro() will be used"
        update_fn=${update_micro}
    elif [[ $1 =~ ^debian ]] || \
            [[ $1 =~ ^ubuntu ]]; then
        inform "Debian - update_deb() will be used"
        update_fn=${update_deb}
    elif [[ $1 =~ ^fedora ]]; then
        inform "Fedora - update_fedora() will be used"
        update_fn=${update_fedora}
    elif [[ $1 =~ ^sles-es ]]; then
        inform "SLES ES - update_res() will be used"
        update_fn=${update_res}
   elif [[ $1 =~ ^centos ]]; then
        inform "Centos - update_res() will be used"
        update_fn=${update_res}
   elif [[ $1 =~ ^sll ]]; then
        inform "SUSE Liberty Linux - update_res() will be used"
        update_fn=${update_res}
   elif [[ $1 =~ ^archlinux ]]; then
        inform "Arch Linux - update_arch() will be used"
        update_fn=${update_arch}
    else
        update_fn=""
    fi


}

update_vm() {
    local update_fn=""

    # get update function for VM
    get_update_function $1

    # if not found, try to remove ${USER}- prefix
    [ -z "${update_fn}" ] && get_update_function ${1#*-}
    if [ -z "${update_fn}" ]; then
        warn "Unknown distribution, I don't know how to update $1."
        return 0
    fi

    # timeout: send SIGKILL if SIGTERM does not work (which sometimes happens)
    call_ssh "$2" "${update_fn}"
    ret=$?

    # if something bad happened, lets report it
    if [ $ret -ne 0 ]; then
        echo "Update of '$1' failed with return code '$ret'."
    fi
    return "$ret"
}

export -f update_vm
export update_9 update_10_11 update_12 update_deb update_res
export update_fedora update_tumbleweed update_micro update_arch

send_report() {
    [ "${error_report}" = true ] || return
    inform "Sending report to ${L3VM_REPORT_EMAIL}"
    echo "$REPORT" | mail -s "$HOSTNAME autoupdate" "${L3VM_REPORT_EMAIL}"
    if [ $? -ne 0 ]; then
        error -l "Problem with sending report"
    fi
    exit
}


# update password on remote VM for root and user "user"
# $1 hostname
update_password() {
    local fn="{ echo root:'$NEWPASS' ; echo user:'$NEWPASS'; } | chpasswd"
    local RET=0
    inform "Setting password to user & root on $1"
    call_ssh "${2}" "${fn}"
    RET=$?
    if [ $RET -ne 0 ]; then
        warn "Changing passwords on $1 failed with return value $RET. "
    fi
    inform "Done"
    return $RET
}

wait_for_command() {
    local i ready
    inform " * $3"
    for i in $(seq 1 "$2"); do
        echo -n '.'
        if eval $1 &> /dev/null; then
            ready=y
            break
        fi
        sleep 1
    done
    if [ "$ready" ]; then
        echo "done"
        return 0
    else
        echo "timeout"
        return 1
    fi
}

wait_for_ssh() {
    wait_for_command \
        "netcat -z '$1' 22" \
        "${L3VM_SSH_START_TIMEOUT}" \
        "Waiting for SSH..."
}

wait_for_off() {
    wait_for_command \
        "[ \"\$(status '$1')\" = \"shut off\" ]" \
        "${L3VM_POWEROFF_TIMEOUT}" \
        "Waiting for machine to be turned off..."
}

process_machine() {
#   $1 virtual machine name

    local vm="$1"
    local broken=false
    local retval=0
    local myret=0

    echo ""
    inform "Processing '$vm'"
    if ! is_machine "$vm"; then
        error "Machine '$vm' is not defined"
        failed[${#failed[@]}]="$vm"
        return 1
    fi

    # OK, machine exists, is it running?
    is_running "$vm"
    local running=$?
    if [ "$running" -eq 0 ] && ! $allow_running ; then
        error "Machine '$vm' is running and it should be not"
        failed[${#failed[@]}]="$vm"
        return 2
    fi

    # if the machine is not running, I need to start it
    if [ "$running" -ne 0 ]; then
        start "$vm"
        if [ $? -ne 0 ]; then
            error "Failed to start '$vm'"
            return 3
        fi
        local started=yes
    fi

    local address="$(get_address -s "$vm")"

    if ! wait_for_ssh "$address"; then
        error "Waiting for SSH reached time out"
        broken=true
        # cannot just exit, we need to shutdown the virtual machine now
    fi

    # do the needful here!
    if ! $broken; then
        if [ ${ssh_keys} = true ]; then
            update_ssh_keys "$vm" "$address"
            myret=$?
            [ $retval -eq 0 ] && retval=$myret
        fi
        if [ ${change_password} = true ]; then
            update_password "$vm" "$address"
            myret=$?
            [ $retval -eq 0 ] && retval=$myret
        fi
        if [ ${update} = true ]; then
            update_vm "$vm" "$address"
            myret=$?
            [ $retval -eq 0 ] && retval=$myret
        fi
    fi

    if [ "$running" -eq 0 ]; then
        # machine was running already, keep it running
        return $retval
    fi

    stop "$vm"
    if ! wait_for_off "$vm"; then
        error -l "Couldn't gracefully shut down the machine."
        if $allow_force_off; then
            inform -l "Stopping virtual machine with force"
            ( OPT_FORCE=1 stop "$vm" )
        else
            return 5
        fi
    fi

    return $retval
}

manage_usage() {
    cat <<EOF
manage: perform action on set of virtual machines using l3vm library.

Usage:
    $(basename "$0") manage < action > [ options ] < virtual_machine > [ ... ]

    Actions:
        Select one or more actions to perform on virtual machines.
        You can set authorized_keys on the target or install sshpass and set
        current password to SSHPASS variable to avoid password prompt.
      --update, -u                update selected virtual machines
      --ssh-keys, -k              copy /etc/l3vm/authorized_keys to given virtual
                                  machines to ~root/.ssh/ and ~user/.ssh/
      --change-password, -p       update password for "root" and "user" accounts.
                                  New password is read from NEWPASS env variable.
    Set of virtual machines:
      --templates, -t             perform the actions on all templates
      --supported-templates, -s   perform the actions only on supported templates
                                  (not listed in L3VM_UPDATE_SKIP array in /etc/l3vm/config.sh)

    Other options
      --error-report, -e          send email report if any error is detected
      --allow-running             perform action even when the virtual machine
                                  is already running
                                  (default: do not mess with running VMs)
      --allow-force-off           when virtual machine doesn't respond to
                                  power button, use the force, Luke
                                  (default: be gentle)
      --stop-on-error             if error occurs, continue with other virtual
                                  machines anyway
                                  (default: do not stop)
      --debug, -d                 show debug information
      --help, -h                  This help

EOF
}

cleanup_usage() {
    cat <<EOF
cleanup: find and remove abandoned virtual machines

Usage:
    $(basename "$0") cleanup [ options ]

    Options:
      --allow-running, -r         Allow to remove running machines (incl. suspended)
      --dryrun, -D                Do not remove VMs, only report them
      --debug, -d                 Show debug information
      --regexp, -x < regexp >     Select only machines whose name match to regular expression
      --no-grace-period, -n       Ignore grace period after registration expiration
      --no-confirm, -c            Destroy virtual machines without confirmation
      --help, -h                  This help

EOF
}


main_usage() {
    cat <<EOF
l3vmadm: perform l3vm administrative tasks

Usage:
    $(basename "$0") [cleanup|manage] [ options ]

    cleanup:
        * find and remove abandoned machines
    manage:
        * connect to virtual machines to update or reconfigure them

EOF
}


#
# configure the run
###################

# We can define the action using external program or exported function

#usage
#exit

manage_run() {
    vms=( )
    origin_args="$@"

    eval set -- $(getopt -o 'ukpetsrldh' --long 'update,ssh-keys,change-password,error-report,templates,supported-templates,allow-running,allow-force-off,stop-on-error,log,debug,help' -n "$(basename "$0")" -- "$@")
    while [ "$1" ]; do
        case "$1" in
            "-r"|"--allow-running")
                allow_running=true
                shift ;;
            "--allow-force-off")
                allow_force_off=true
                shift ;;
            "--stop-on-error")
                stop_on_error=true
                shift ;;
            "-u"|"--update")
                update=true
                shift ;;
            "-k"|"--ssh-keys")
                ssh_keys=true
                shift ;;
            "-p"|"--change-password")
                change_password=true
                shift ;;
            "-e"|"--error-report")
                error_report=true
                shift ;;
            "-t"|"--templates")
                vms=("${L3VM_TEMPLATES[@]}")
                shift ;;
            "-s"|"--supported-templates")
                vms=($(list_supported_templates))
                shift ;;
            "-l"|"--log")
                # redirect output to log file
                exec 1>>"${L3VM_LOG_FILE}"
                exec 2>&1
                shift ;;
            "-d"|"--debug")
                inform "debug mode enabled"
                set -x
                L3VM_DEBUG=true
                shift ;;
            "-h"|"--help")
                manage_usage
                exit ;;
            "--")
                shift
                break ;;
        esac
    done

    inform ""
    inform "===================================================="
    inform " $(basename $0) is going to manage machines."
    inform " input: ${origin_args}"
    inform " date: $(date)"
    inform "===================================================="

    if [ "$update" = false -a "$ssh_keys" = false -a "$change_password" = false  ]; then
        error "No action defined."
        exit 7
    fi

    if ! [ -z "$*" ]; then
        for i in "$@"; do
            vms_exp=" ${vms[*]} "
            if ! [[ ${vms_exp} =~ $i ]]; then
                vms[${#vms[@]}]="$i"
            fi
        done
    fi

    if [ ${#vms[@]} -eq 0 ]; then
        error "No virtual machines selected."
        exit 6
    fi

    if [ ${change_password} = true -a -z "${NEWPASS}" ]; then
        error "New password must be specified in NEWPASS env variable."
        exit 8
    fi

    SSH_EXTRA_ARGS=""
    if [ "${L3VM_DEBUG}" = true ]; then
        SSH_EXTRA_ARGS="${SSH_EXTRA_ARGS} -vvvv"
        SSH_CMD_TRACE=1
    fi

    inform -l "VMs: " "${vms[@]}"
    local operations="$([ "$change_password" = true ] && echo " change-password")\
$([ "$update" = true ] && echo " update")$([ "$ssh_keys" = true ] && echo " ssh-keys")"
    inform -l "Operations:$operations"

    global_retval=0

    REPORT="
Date: $(date)
Arguments: $@
Operations: $operations
Machines: ${vms[@]}

"

    for machine in "${vms[@]}"; do
        process_machine "$machine" > >(tee -a "${L3VM_REPORT_TMP}") 2> >(tee -a "${L3VM_REPORT_TMP}" >&2)
        ret=$?

        if [ $ret -ne 0 ]; then
            global_retval=${retval}
            REPORT+="  error $ret      ${machine}"$'\n'
        else
            REPORT+="  ok              ${machine}"$'\n'
        fi
    done

    if ! [ "${global_retval}" = 0 ]; then
        inform "===================================================="
        warn -l "Update went wrong. Finished on $(date)."
        inform "===================================================="
        inform ""
        REPORT+=""$'\n'
        REPORT+="Update went wrong. Finished on $(date)."$'\n'
        REPORT+=""$'\n'
        REPORT+="Details:"$'\n'
        # strip ANSII colors before sending
        REPORT+="$(sed $'s/\033\\[[0-9;]*m//g' <"${L3VM_REPORT_TMP}")"$'\n'
        send_report
    else
        inform "===================================================="
        inform -l "Update went smoothly. Finished on $(date)."
        inform "===================================================="
        inform ""
    fi

    # return last known error or 0 on success
    return ${global_retval}
}

cleanup_run() {
    dryrun=false
    allow_running=false
    vm_regexp=""
    allow_grace_period=true
    no_confirm=false
    origin_args="$@"

    eval set -- $(getopt -o 'cdDhx:rnl' --long 'no-confirm,debug,dryrun,help,regexp:,allow-running,no-grace-period,log' -n "$(basename "$0")" -- "$@")
    while [ "$1" ]; do
        case "$1" in
            "-c"|"--no-confirm")
                no_confirm=true
                shift ;;
            "-r"|"--allow-running")
                allow_running=true
                shift ;;
            "-x"|"--regexp")
                shift
                vm_regexp="${1}"
                shift ;;
            "-d"|"--debug")
                inform "debug mode enabled"
                set -x
                L3VM_DEBUG=true
                shift ;;
            "-D"|"--dryrun")
                dryrun=true
                shift ;;
            "-n"|"--no-grace-period")
                allow_grace_period=false
                shift ;;
            "-h"|"--help")
                cleanup_usage
                return 0;;
            "-l"|"--log")
                # redirect output to log file
                exec 1>>"${L3VM_LOG_FILE}"
                exec 2>&1
                shift ;;
            --)
                shift
                break ;;
            *)
                echo "Unknown parameter '$1' found. Exiting..."
                exit 1
        esac
    done

    inform ""
    inform "===================================================="
    inform " $(basename $0) is going to find abandoned machines."
    inform " input: ${origin_args}"
    inform " date: $(date)"
    inform "===================================================="

    vms=( )
    list_clones > /dev/null
    for machine in "${CLONES[@]}"; do
        if ! is_shutoff "${machine}" && [ ${allow_running} = false ]; then
            continue
        fi
        if [ -n "${vm_regexp}" ] && ! [[ ${machine} =~ ${vm_regexp} ]]; then
            continue
        fi
        get_reservation_status "${machine}"
        if [ ${allow_grace_period} = true ]; then
            # still before reservation + grace period
            if [ "${now}" -lt "${abandoned}" ]; then
                continue
            fi
        else
            # still reserved
            if [ "${now}" -lt "${reservation}" ]; then
                continue
            fi
        fi

        vms[${#vms[@]}]="$machine"
    done

    if [ ${#vms[@]} -eq 0 ]; then
        inform "No abandoned machine found."
        return 0
    fi

    inform "Found:" "${vms[@]}"

    if [ ${dryrun} = true ]; then
        inform "Dry run only, returning without deleting any virtual machine."
        return 0
    fi

    inform -l "$(basename $0) cleanup is going to remove selected machines:" "${vms[@]}."
    for machine in "${vms[@]}"; do
        if [ ${no_confirm} = true ]; then
            l3vm destroy -f ${machine}
        else
            l3vm destroy ${machine}
        fi
    done

    return 0
}


RET=0
case "$1" in
    --help|-h|help)
        main_usage ;;
    "")
        main_usage ;;
    manage)
        shift
        manage_run "$@"
        RET=$?
        ;;
    cleanup)
        shift
        cleanup_run "$@"
        RET=$?
        ;;
    *)
        error "Unknown command" "$1"
        RET=1
        ;;
esac

exit $RET

