#!/bin/bash

#set -x

# Cross Pool Migration and incremental replication Tool for Proxmox VMs using Ceph.
# Author: Bastian Mäuser <bma@netz.org>        

declare -r VERSION=0.3
declare -r NAME=$(basename "$0")
declare -r PROGNAME=${NAME%.*}

declare -r PVE_DIR="/etc/pve"
declare -r PVE_NODES="$PVE_DIR/nodes"
declare -r QEMU='qemu-server'
declare -r QEMU_CONF_CLUSTER="$PVE_NODES/*/$QEMU"
declare -r EXT_CONF='.conf'
declare -r PVFORMAT_FULL='elapsed:%t remaining:%e current:%r average:%a %b %p'
declare -r PVFORMAT_SNAP='elapsed:%t current:%r average:%a %b'

declare -r LOG_FILE=$(mktemp)

declare -A -g pvnode
declare -A -g dstpvnode
declare -A -g svmids
declare -A -g dvmids

declare opt_destination
declare opt_vm_ids=''
declare opt_snapshot_prefix='mirror-'
declare opt_rewrite=''
declare -i opt_prefix_id
declare opt_exclude_vmids=''
declare -i opt_debug=0
declare -i opt_dry_run=0
declare -i opt_syslog=0
declare -i opt_lock=1
declare -i opt_keepslock=0
declare -i opt_keepdlock=0
declare -i opt_overwrite=0
declare -i opt_online=0
declare -i opt_keep_local=0
declare -i opt_keep_remote=0

declare -r redstconf='^\/etc\/pve\/nodes\/(.*)\/qemu-server\/([0-9]+).conf$'
declare -r recephimg='([a-zA-Z0-9]+)\:(.*)'

function usage(){
   shift

    if [ "$1" != "--no-logo" ]; then
        cat << EOF
 _____                                 
|     |___ ___ ___ ___ ___ _ _ ___ ___ 
|   --|  _| . |_ -|_ -| . | | | -_|  _|
|_____|_| |___|___|___|___|\_/|___|_|  

EOF
    fi

    cat << EOF
Mirror tool Proxmox VMs on Ceph storages

Usage:
    $PROGNAME <COMMAND> [ARGS] [OPTIONS]
    $PROGNAME help
    $PROGNAME version
 
    $PROGNAME mirror     --vmid=<string> --destination=<destionationhost> --pool=<targetpool> --keeplocal=n --keepremote=n
Commands:
    version              Show version program
    help                 Show help program
    mirror               Replicate a stopped VM to another Cluster (full clone)

Options:
    --vmid               The source+target ID of the VM/CT, comma separated (eg. --vmid=100:100,101:101), or all for all
    --prefixid           Prefix for VMID's on target System [optional]
    --excludevmids       Exclusde VM IDs when using --vmid==all
    --destination        Target PVE Host in target pool. e.g. --destination=pve04
    --pool               Ceph pool name in target pool. e.g. --pool=data
    --keeplocal          How many additional Snapshots to keep locally. e.g. --keeplocal=2
    --keepremote         How many additional Snapshots to keep remote. e.g. --keepremote=2
    --rewrite            PCRE Regex to rewrite the Config Files (eg. --rewrite='s/(net0:)(.*)tag=([0-9]+)/\1\2tag=1/g' would
                         change the VLAN tag from 5 to 1 for net0.
Switches:    
    --online             Allow online Copy
    --nolock             Don't lock source VM on Transfer (mainly for test purposes)
    --keep-slock         Keep source VM locked on Transfer
    --keep-dlock         Keep VM locked after transfer on Destination
    --overwrite          Overwrite Destination
    --debug              Show Debug Output

Report bugs to <mephisto@mephis.to>

EOF
    exit 1
}

function parse_opts(){
#    local action=$1
    shift

    local args
    args=$(getopt \
           --options '' \
           --longoptions=vmid:,prefixid:,excludevmids:,destination:,pool:,keeplocal:,keepremote:,rewrite:,online,nolock,keep-slock,keep-dlock,overwrite,dry-run,debug \
           --name "$PROGNAME" \
           -- "$@") \
           || end_process 128

    eval set -- "$args"

    while true; do    
      case "$1" in
        --vmid) opt_vm_ids=$2; shift 2;;
        --prefixid) opt_prefix_id=$2; shift 2;;
        --excludevmids) opt_exclude_vmids=$2; shift 2;;
        --destination) opt_destination=$2; shift 2;;
        --pool) opt_pool=$2; shift 2;;
        --keeplocal) opt_keep_local=$2; shift 2;;
        --keepremote) opt_keep_remote=$2; shift 2;;
        --rewrite) opt_rewrite=$2; shift 2;;
        --online) opt_online=1; shift 2;;
        --dry-run) opt_dry_run=1; shift;;
        --debug) opt_debug=1; shift;;
        --nolock) opt_lock=0; shift;;
        --keep-slock) opt_keepslock=1; shift;;
        --keep-dlock) opt_keepdlock=1; shift;;
        --overwrite) opt_overwrite=1; shift;;
        --) shift; break;;
        *) break;;
      esac
    done

    if [ $opt_debug -eq 1 ]; then
        log info "============================================"
        log info "Proxmox Crosspool Migration: $VERSION";
        log info "============================================"
        log info "Proxmox VE Version:"

        pveversion

        log info "============================================"
    fi

    [ -z "$opt_vm_ids" ] && { log info "VM id is not set."; end_process 1; }

    if [ "$opt_vm_ids" = "all" ]; then
        local all=''
        local data=''
        local cnt=''

        all=$(get_vm_ids "$QEMU_CONF_CLUSTER/*$EXT_CONF" "$LXC_CONF_CLUSTER/*$EXT_CONF")
        all=$(echo "$all" | tr ',' "\n")
        opt_exclude_vmids=$(echo "$opt_exclude_vmids" | tr ',' "\n")
        for id in $all; do
            cnt=$(echo $opt_exclude_vmids | grep -o $id|wc -w)
            if [ $cnt == 0 ]; then
                vm_ids=$(echo "$vm_ids$id:$opt_prefix_id$id,")
            fi
        done
        vm_ids=$(echo "$vm_ids" | tr ',' "\n")
    else
        vm_ids=$(echo "$opt_vm_ids" | tr ',' "\n")
    fi
}

function map_vmids_to_host(){
    for node in $(/usr/bin/pvecm nodes | tail +5 | tr -s ' ' | cut -d' ' -f 4)
    do
        for vm in $(ssh root@$node qm list|tail +2|tr -s ' '|cut -f 2 -d' ')
        do
            pvnode[$vm]=$node
        done
    done
}

function map_vmids_to_dsthost(){
    for node in $(ssh $1 /usr/bin/pvecm nodes | tail +5 | tr -s ' ' | cut -d' ' -f 4)
    do
        for vm in $(ssh root@$node qm list|tail +2|tr -s ' '|cut -f 2 -d' ')
        do
            dstpvnode[$vm]=$node
        done
    done
}

function exist_file(){
    local file=''
    for file in $1; do
        [ -e "$file" ] && return 0 || return 1
        break
    done
}

function lookupcephpool() {
    pvehost=$1
    pvepoolname=$2
    res=$(ssh $pvehost cat /etc/pve/storage.cfg | sed -n "/rbd: $pvepoolname/,/^$/p" | grep pool | cut -d " " -f 2)
    echo $res
}

function get_vm_ids(){
    local data=''
    local conf=''

    while [ $# -gt 0 ]; do
        for conf in $1; do
            [ ! -e "$conf" ] && break

            conf=$(basename "$conf")
            [ "$data" != '' ] && data="$data,"
            data="$data${conf%.*}"
        done
        shift
    done

    echo "$data"
}

function get_config_file(){
    local file_config=''

    if exist_file "$QEMU_CONF_CLUSTER/$vm_id$EXT_CONF"; then
        file_config=$(ls $QEMU_CONF_CLUSTER/$vm_id$EXT_CONF)

    else
        log error "VM $vm_id - Unknown technology or VMID not found: $QEMU_CONF_CLUSTER/$vm_id$EXT_CONF"
        end_process 128
    fi 

    echo "$file_config"
}

function get_disks_from_config(){
    local disks;
    local file_config=$1

    #disks available for vm/ct
    #exclude no backup
    #read current config
    disks=$(while read -r line; do
                [[ "$line" == "" ]] && break
                echo "$line" 
            done < "$file_config" | \
            grep -P '^(?:((?:virtio|ide|scsi|sata|mp)\d+)|rootfs): ' | \
            grep -v -P 'cdrom|none' | \
            grep -v -P 'backup=0' | \
            awk '{ split($0,a,","); split(a[1],b," "); print b[2]}')

    echo "$disks"
}

function log(){
    local level=$1
    shift 1
    local message=$*

    case $level in
        debug) 
            if [ $opt_debug -eq 1 ]; then
                echo -e "$(date "+%F %T") DEBUG: $message";
                echo -e "$(date "+%F %T") DEBUG: $message" >> "$LOG_FILE";
            fi    
            ;;

        info) 
            echo -e "$message"; 
            echo -e "$message" >> "$LOG_FILE";
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" "$message"
            ;;

        warn)
            echo "WARNING: $message" 1>&2
            echo -e "$message" >> "$LOG_FILE";            
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" -p daemon.warn "$message"
            ;;
        
        error)
            echo "ERROR: $message" 1>&2
            echo -e "$message" >> "$LOG_FILE";            
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" -p daemon.err "$message"
            ;;

        *)  
            echo "$message" 1>&2
            echo -e "$message" >> "$LOG_FILE";            
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" "$message"
            ;;
    esac
}

function map_source_to_destination_vmid(){
    for vm_id_pair in $vm_ids; do
        svid=$(echo "$vm_id_pair" | sed -r -e "s/^([0-9]+):([0-9]+)$/\1\n/")
        dvid=$(echo "$vm_id_pair" | sed -r -e "s/^([0-9]+):([0-9]+)$/\2\n/")
        svmids="$svmids"$'\n'"$svid"
        dvmids[$svid]=$dvid
    done
}

function mirror() {
    local -i rc=0;
    parse_opts "$@"

    local timestamp; timestamp=$(date +%Y%m%d%H%M%S)

    log info "ACTION: Onlinemirror"
    log info "Start mirror $(date "+%F %T")"

    #create pid file
    local pid_file="/var/run/$PROGNAME.pid"
    if [[ -e "$pid_file" ]]; then
        local pid; pid=$(cat "${pid_file}")
        if ps -p "$pid" > /dev/null 2>&1; then
          log error "Process already running with pid ${pid}"
          end_process 1
        fi
    fi
    if ! echo $$ > "$pid_file"; then
        log error "Could not create PID file $pid_file"
        end_process 1
    fi

    map_source_to_destination_vmid
    map_vmids_to_host
    map_vmids_to_dsthost "$opt_destination"

    for vm_id in $svmids; do
        local file_config; file_config=$(get_config_file)
        [ -z "$file_config" ] && continue
        local disk=''
        dvmid=${dvmids[$vm_id]}


        srcvmgenid=$(cat $PVE_NODES/"${pvnode[$vm_id]}"/$QEMU/"$vm_id".conf|grep vmgenid|sed -r -e 's/^vmgenid:\s(.*)/\1/')
        dstvmgenid=$(ssh $opt_destination cat $PVE_NODES/"${dstpvnode[$dvmid]}"/$QEMU/"$dvmid".conf 2>/dev/null|grep vmgenid|sed -r -e 's/^vmgenid:\s(.*)/\1/')
        log debug "Checking for VM $dvmid on Destination Host $opt_destination $QEMU_CONF_CLUSTER"
        log debug "DVMID: $dvmid"
        conf_on_destination=$(ssh $opt_destination "ls -d $QEMU_CONF_CLUSTER/$dvmid$EXT_CONF 2>/dev/null")
        [[ "$conf_on_destination" =~ $redstconf ]]
        host_on_destination=${BASH_REMATCH[1]}

        if [ $host_on_destination ]; then
            dststatus=$(ssh root@${dstpvnode[$dvmid]} qm status $dvmid|cut -d' ' -f 2)
            if [ $dststatus == "running" ]; then
                log error "Destination VM is running. bailing out"
                end_process 255
            fi
        fi
        #Check if source VM is running
        srcstatus=$(ssh root@${pvnode[$vm_id]} qm status $vm_id|cut -d' ' -f 2)
        if [ $srcstatus == "running" ] && [ $opt_online -eq 0 ]; then
            log error "Source VM is running .. exiting"
            end_process 1
        fi

        if [ ! "$host_on_destination" ] || [ $opt_overwrite -eq 1 ]; then
            if [ "$host_on_destination" ] && [ $srcvmgenid != $dstvmgenid ]; then
                log error "Source VM genid ($srcvmgenid) doesn't match destination VM genid ($dstvmgenid). This should not happen. Bailing out.."
                end_process 255
            fi
            log info "VM $vm_id - Transmitting Config for to destination $opt_destination VMID $dvmid"
            rewriteconfig $PVE_NODES/"${pvnode[$vm_id]}"/$QEMU/"$vm_id".conf $opt_destination "$opt_pool" $PVE_NODES/"$opt_destination"/$QEMU/"$dvmid".conf "$dvmid"
            map_vmids_to_dsthost "$opt_destination"
        fi

        #Lock on source + destination
        if [ $opt_lock -eq 1 ]; then
            ssh root@"${pvnode[$vm_id]}" qm set "$vm_id" --lock backup
        fi

        if [ $opt_lock -eq 1 ]; then
            ssh root@"${dstpvnode[$dvmid]}" qm set "$dvmid" --lock backup
        fi

        #Freeze, take Rbd Snapshot then unfreeze
        vm_freeze "$vm_id" "${pvnode[$vm_id]}"
        for disk in $(get_disks_from_config "$file_config"); do
            src_image_spec=$(get_image_spec "$disk")
            create_snapshot "$src_image_spec@$opt_snapshot_prefix$timestamp"
        done
        vm_unfreeze "$vm_id" "${pvnode[$vm_id]}"

        for disk in $(get_disks_from_config "$file_config"); do
            log debug "VMID: $vm_id Disk: $disk DESTVMID: $dvmid"
            src_image_spec=$(get_image_spec "$disk")
            [ -z "$src_image_spec" ] && continue
            dst_image_spec=$(echo $src_image_spec | sed -r -e "s/(.*\/[a-zA-Z0-9]+\-)([0-9]+)(\-[a-zA-Z0-9]+\-[0-9]+)/\1$dvmid\3/")
            [ -z "$dst_image_spec" ] && continue
            [[ $disk =~ $recephimg ]]
	        src_image_pool_pve=${BASH_REMATCH[1]}
            src_image_pool=$(lookupcephpool "localhost" ${BASH_REMATCH[1]})
            src_image_name=${BASH_REMATCH[2]}
            [[ $dst_image_spec =~ ^.*\/(.*)$ ]]
            dst_image_name=${BASH_REMATCH[1]}-$src_image_pool_pve
            dst_image_pool=$(lookupcephpool $opt_destination $opt_pool)
            snapshot_name="@$opt_snapshot_prefix$timestamp"
            localsnapcount=$(rbd ls -l $src_image_pool | grep $src_image_name@$opt_snapshot_prefix | cut -d ' ' -f 1|wc -l)
            if [ $localsnapcount -ge 2 ]; then
            # we have at least 2 local snapshots, to we can make an incremental copy
                currentlocal=$(rbd ls -l $src_image_pool | grep $src_image_name@$opt_snapshot_prefix | cut -d ' ' -f 1|tail -n 1)
                localts=$(rbd ls -l $src_image_pool | grep $src_image_name@$opt_snapshot_prefix | cut -d ' ' -f 1 | sed -r -e 's/.*@mirror-(.*)/\1/')
            fi
            latestremote=$(ssh $opt_destination rbd ls -l $dst_image_pool | grep $dst_image_name@$opt_snapshot_prefix | cut -d ' ' -f 1|tail -n 1)
            if [ $latestremote ]; then
                [[ $latestremote =~ ^.*@$opt_snapshot_prefix([0-9]+)$ ]]
                latestremotets=${BASH_REMATCH[1]}
                for ts in $localts; do
                    if [ $ts == $latestremotets ]; then
                        basets=$ts
                    fi
                done
            fi
            if [ -z $basets ]; then
                log debug "No matching Snapshot found on destination - Full Copy $src_image_pool/$src_image_name$snapshot_name to $dst_image_pool/$dst_image_name"
                #snapts=$(echo $currentlocal | sed -r -e 's/.*@mirror-(.*)/\1/')
                snapshotsize=$(rbd du --pretty-format --format json $src_image_pool/$src_image_name|jq '.images[] | select (.snapshot_id == null) | {provisioned_size}.provisioned_size'|tail -1)
                log debug "snapsize: $snapshotsize "
                xmitjob="rbd export --rbd-concurrent-management-ops 8 $src_image_pool/$src_image_name$snapshot_name --no-progress - | pv -s $snapshotsize -F \"VM $vm_id - Full xmit: $PVFORMAT_FULL\" | ssh $opt_destination rbd import --image-format 2 - $dst_image_pool/$dst_image_name"
                # create initial snapshot on destination
                if ! do_run "$xmitjob"; then
                    log error "Transmitting Image failed"
                    return 1
                fi
                cmd="ssh $opt_destination rbd snap create $dst_image_pool/$dst_image_name$snapshot_name"
                do_run $cmd
            else
                log debug "Basecopy + snapshot on destination - let's just transfer the diff"
                log debug "sizer: rbd diff $src_image_pool/$currentlocal --from-snap $opt_snapshot_prefix$basets|gawk --bignum '{ SUM += \$2 } END { print SUM }'"
                snapshotsize=$(rbd diff $src_image_pool/$currentlocal --from-snap $opt_snapshot_prefix$basets|gawk --bignum '{ SUM += $2 } END { print SUM }')
                log debug "snapshotsize: $snapshotsize"
                if [ -z "$snapshotsize" ]; then
                    #disk was not attached, or really nothing has changed..
                    snapshotsize=0
                fi
                xmitjob="rbd export-diff --no-progress --from-snap $opt_snapshot_prefix$basets $src_image_pool/$currentlocal - | pv -F \"VM $vm_id - Snap xmit: $PVFORMAT_SNAP\" | ssh $opt_destination rbd import-diff --no-progress - $dst_image_pool/$dst_image_name"
                log debug "xmitjob: $xmitjob"
                if ! do_run "$xmitjob"; then
                    log error "Transmitting Image failed"
                    return 1
                fi
                do_housekeeping "localhost" "$src_image_pool" "$src_image_name" $opt_keep_local
                do_housekeeping "$opt_destination" "$dst_image_pool" "$dst_image_name" $opt_keep_remote
            fi
            unset basets
        done
        if [ ! $opt_keepslock -eq 1 ]; then
            ssh root@${pvnode[$vm_id]} qm unlock $vm_id
            log info "VM $vm_id - Unlocking source VM $vm_id"
        fi
        if [ $opt_keepdlock -eq 0 ]; then
            ssh root@${dstpvnode[$dvmid]} qm unlock $dvmid
            log info "VM $dvmid - Unlocking destination VM $dvmid"
        fi
    done
    log info "Finnished mirror $(date "+%F %T")"
}

function do_housekeeping(){
    horst=$1
    rbdpool=$2
    rbdimage=$3
    keep=$4
    snapshotstokill=$(ssh $horst rbd ls -l $rbdpool | grep $rbdimage@$opt_snapshot_prefix | cut -d ' ' -f 1|head -n -1 |head -n -$keep)
    log info "Houskeeping $horst $rbdpool $rbdimage, keeping previous $keep Snapshots"
    for snap in $snapshotstokill; do
        cmd="ssh $horst rbd snap rm $rbdpool/$snap"
        if ! do_run $cmd; then
            log error "Housekeeping failed: $cmd"
            return 1
        fi
    done
}

function create_snapshot(){
    local snap="$1"
    log info "VM $vm_id - Creating snapshot $snap"
    if ! do_run "rbd snap create $snap"; then    
        return 1;
    fi
}

function vm_freeze() {
    local fvm=$1;
    local fhost=$2;
    status=$(ssh root@$fhost qm status $fvm|cut -d' ' -f 2)
    if ! [[ "$status" == "running" ]]; then
        log info "VM $fvm - Not running, skipping fsfreeze-freeze"
        return
    fi
    local cmd="ssh root@$fhost /usr/sbin/qm guest cmd $fvm fsfreeze-freeze"
    log info "VM $fvm - Issuing fsfreeze-freeze to $fvm on $fhost"
    do_run "$cmd"
    rc=$?
    log debug "vm_freeze() return $rc"
}

function vm_unfreeze() {
    local fvm=$1;
    local fhost=$2;
    status=$(ssh root@$fhost qm status $fvm|cut -d' ' -f 2)
    if ! [[ "$status" == "running" ]]; then
        log info "VM $fvm - Not running, skipping fsfreeze-thaw"
        return
    fi
    local cmd="ssh root@$fhost /usr/sbin/qm guest cmd $fvm fsfreeze-thaw"
    log info "VM $fvm - Issuing fsfreeze-thaw to $fvm on $fhost"
    do_run "$cmd"
    rc=$?
    log debug "vm_unfreeze() return $rc"
}

function rewriteconfig(){
    local oldconfig=$1
    local dst=$2
    local newpool=$3
    local newconfig=$4
    local newvmid=$5
    local sedcmd
    if [ ! -z $opt_rewrite ]; then
        sedcmd='sed -r -e '$opt_rewrite
    else
        sedcmd='sed -e /^$/,$d'
    fi
    cat "$oldconfig" | sed -r -e "s/^(virtio|ide|scsi|sata|mp)([0-9]+):\s([a-zA-Z0-9]+):(.*)-([0-9]+)-disk-([0-9]+).*,(.*)$/\1\2: $newpool:\4-$newvmid-disk-\6-\3,\7/g" | $sedcmd | sed -e '/^$/,$d' | grep -v "^parent:\s.*$" | ssh $dst "cat - >$newconfig"
}

function checkvmid(){
    local dst=$1
    local vmid=$2
    cmd="ssh $dst ls -l $QEMU_CONF_CLUSTER/$vmid.conf|wc -l"
    rval=$($cmd)
    echo $rval
}

function do_run(){
    local cmd=$*;
    local -i rc=0;

    if [ $opt_dry_run -eq 1 ]; then
        echo "$cmd"
        rc=$?
    else
        log debug "$cmd"
        eval "$cmd"
        rc=$?
        [ $rc != 0 ] && log error "$cmd"
        log debug "return $rc ps ${PIPESTATUS[@]}"
    fi

    return $rc
}

function end_process(){
    local -i rc=$1;
#    if ! [[ -z "$startts" && -z "$endts" ]]; then
#        local -i runtime=$(expr $endts - $startts)
#        local -i bps=$(expr $bytecount/$runtime)
#    fi
#    local subject="Ceph [VM:$vmok/$vmtotal SS:$snapshotok/$snapshottotal EX:$exportok/$exporttotal] [$(bytesToHuman "$bytecount")@$(bytesToHuman "$bps")/s]"
#    [ $rc != 0 ] && subject="$subject [ERROR]"

    #send email
#    local mail;
#    local mailhead="Backup $imgcount Images in $vmcount VMs (Bytes: $bytecount)"
#    for mail in $(echo "$opt_addr_mail" | tr "," "\n"); do
#        do_run "cat '$LOG_FILE' | mail -s '$subject' '$mail'"
#    done

    #remove log
#    rm "$LOG_FILE"

    exit "$rc";
}

function get_image_spec(){
    local image_spec;
    local disk="$1"
    
    #if krbd enable
    image_spec=$(pvesm path "$disk" | grep '^/dev/rbd/' | sed -e "s/^\/dev\/rbd\///")
    if [ -z "$image_spec" ]; then
        image_spec=$(pvesm path "$disk" | grep '/ceph/' | awk '{ split($0,a,":"); print a[2]}')
    fi

    echo "$image_spec"
}

function main(){    
    [ $# = 0 ] && usage;

    #command
    case "$1" in
        version) echo "$VERSION";;
        help) usage "$@";;
        mirror) mirror "$@";; 
        *) usage;;
    esac

    exit 0;
}

main "$@"