#!/bin/bash
# Original script https://src.openvz.org/scm/ovzl/ovztransfer.git

# Additional ssh opts, another key location for example
#SSH_OPTS="-i /root/id_rsa_target"
SSH_OPTS=""

declare -a VEID_LIST
declare -a TARGET_VEID_LIST
declare -a PIDS_LIST
MIGRATION_STARTED=0

function stop_scripts() {
	[ $MIGRATION_STARTED -eq 0 ] && return
	echo "Stopping all migration processes..."
	kill -HUP -$$
}

trap stop_scripts EXIT

function error() {
    echo $$: $*
    exit 1
}

function usage() {
    echo "Usage: $0 HOSTNAME SOURCE_VEID0:[TARGET_VEID0] ... [SOURCE_VEIDn:[TARGET_VEIDn]]"
    exit 0
}

function migrate() {
    local veid=$1
    local target_veid=$2
    local target=$3
    local ssh_port=$4
    local ssh_opts=$5
    local tmpdir
    local quota_restore_command
    local dir
    local param
    local required_space
    local vefstype=5 #ploop
    local inodes_coeff=8 #1/8
    local ostemplate_rpm
    local pid
    local mult
    local xattrs=""
    local block_pid

    # Start Container
    ssh -p $ssh_port $ssh_opts $target vzctl start $veid > /dev/null 2>&1

    echo "Container $veid: Shutting down all possible services..."

    # Stop all possible processes inside Container

    PS_IGNORE_REGEX="kthread\|khelper\|/"
    PS_TIMEOUT=90
    echo "Waiting for container processes to stop..."
    for pid in `ssh -p $ssh_port $ssh_opts $target vzctl exec $veid ps | grep -v ${PS_IGNORE_REGEX} | awk '!/^ *(PID|1) / {print $1}'`; do
        vzctl exec $veid kill $pid > /dev/null 2>&1
    done

    for i in `seq 1 $PS_TIMEOUT`; do
        remaining=`ssh -p $ssh_port $ssh_opts $target vzctl exec $veid ps | grep -v ${PS_IGNORE_REGEX} | awk '!/^ *(PID|1) / {print $1}'`
        [ "x$remaining" == "x" ] && break
        sleep 1;
    done

    for pid in `ssh -p $ssh_port $ssh_opts $target vzctl exec $veid ps | grep -E -v "^ *(PID|1) " | awk '{print $1}'`; do
        ssh -p $ssh_port $ssh_opts $target vzctl exec $veid kill -9 $pid
    done

    # Dump old quota
    quota_restore_command=`ssh -p $ssh_port $ssh_opts $target /usr/sbin/vzdqdump $veid -f -G -U -T | grep ^ugid: | awk '{
        if ($3 == "1")
                gparm="-g"
        else
                gparm="-u"
        printf("/usr/sbin/setquota %s %s %s %s %s %s -a; ", gparm, $2, $4, $5, $6, $7)
}'`

    # Bind mount
    printf "Source %s: ssh -p $ssh_port $ssh_opts $target vzctl --quiet exec $veid mktemp -d /tmp/bindmnt_XXXXXX\n", "$target"
    tmpdir=`ssh -p $ssh_port $ssh_opts $target vzctl --quiet exec $veid mktemp -d /tmp/bindmnt_XXXXXX`

    printf "Source %s: ssh -p $ssh_port $ssh_opts $target vzctl exec $veid "mount -o bind / $tmpdir; tail -f /dev/null > $tmpdir/tmp/lock"\n", "$target"
    ssh -p $ssh_port $ssh_opts $target vzctl exec $veid "mount -o bind / $tmpdir; tail -f /dev/null > $tmpdir/tmp/lock"
    block_pid=$!

    echo "Container $veid: Copying data..."

    # Calculate needed diskspace
    eval `ssh -p $ssh_port $ssh_opts $target grep "^DISKSPACE=" $VECONFDIR/$veid.conf`
    required_space=`echo $DISKSPACE | sed "s,.*:,,g"`
    # Parse suffix
    required_space_suffix=${required_space:$((${#required_space}-1))}
    case "$required_space_suffix" in
        T)
            required_space=`echo $required_space | sed "s,T$,,g"`
            mult=$((1024*1024*1024))
            ;;
        G)
            required_space=`echo $required_space | sed "s,G$,,g"`
            mult=$((1024*1024))
            ;;
        M)
            required_space=`echo $required_space | sed "s,M$,,g"`
            mult=1024
            required_space=$((required_space*1024))
            ;;
        *)
            mult=1
            ;;
    esac
    # dots
    required_space=`echo $required_space | sed "s,\.,\,,g"`
    required_space=$((required_space*mult))
    # Reserve inodes_coeff for inodes
    required_space=$((required_space + required_space/inodes_coeff))

    # Create destination ploop
    printf "Destination: mkdir -p /vz/private/$target_veid/root.hdd\n"
    mkdir -p /vz/private/$target_veid/root.hdd
    [ $? -ne 0 ] && error "Failed to create root.hdd on $target"
    ploop init -t ext4 -s ${required_space}K /vz/private/$target_veid/root.hdd/root.hds
    [ $? -ne 0 ] && error "Failed to create ploop on $target"

    # Mount ploop
    printf "Destination: ploop mount -m /vz/root/$target_veid /vz/private/$target_veid/root.hdd/DiskDescriptor.xml\n"
    ploop mount -m /vz/root/$target_veid /vz/private/$target_veid/root.hdd/DiskDescriptor.xml
    [ $? -ne 0 ] && error "Failed to mount ploop on $target"

    # Copy data
    # Check for xattrs
    printf "Source %s: ssh -p $ssh_port $ssh_opts $target "vzctl --quiet exec $veid tar --help" | grep "no-xattrs"\n", "$target"
    ssh -p $ssh_port $ssh_opts $target "vzctl --quiet exec $veid tar --help" | grep "no-xattrs"
    [ $? -eq 0 ] && xattrs="--xattrs"

    printf "Destination: "vzctl --quiet exec $veid tar --numeric-owner $xattrs -cz -C $tmpdir ./" | tar --numeric-owner $xattrs -xz -C /vz/root/$target_veid\n"
    ssh -p $ssh_port $ssh_opts $target "vzctl --quiet exec $veid tar --numeric-owner $xattrs -cz -C $tmpdir ./" | tar --numeric-owner $xattrs -xz -C /vz/root/$target_veid
    [ $? -ne 0 ] && error "Failed to copy data"

    # Leave block
    printf "Destination: kill $block_pid\n"
    kill $block_pid

    # Umount target ploop

    printf "Destination: ploop umount /vz/private/$target_veid/root.hdd/DiskDescriptor.xml\n"
    ploop umount /vz/private/$target_veid/root.hdd/DiskDescriptor.xml
    [ $? -ne 0 ] && error "Failed to umount ploop on $target"

    # Fill private area
    printf "Destination: ln -s $vefstype /vz/private/$target_veid/.ve.layout\n"
    ln -s $vefstype /vz/private/$target_veid/.ve.layout
    [ $? -ne 0 ] && error "Failed to create layout on $target"
    for dir in scripts dump fs root.hdd/templates; do
        printf "Destination: mkdir /vz/private/$target_veid/$dir\n"
        mkdir /vz/private/$target_veid/$dir
        [ $? -ne 0 ] && error "Failed to create $dir dir on $target"
    done
    printf "Destination: ln -s root.hdd/templates /vz/private/$target_veid/templates\n"
    ln -s root.hdd/templates /vz/private/$target_veid/templates
    [ $? -ne 0 ] && error "Failed to templates dir link on $target"
    echo -n `hostname` > /vz/private/$target_veid/.owner
    [ $? -ne 0 ] && error "Failed to create owner file on $target"

    # Copy config
    printf "Source %s: scp $ssh_opts $target:$VECONFDIR/$veid.conf /vz/private/$target_veid/ve.conf\n", "$target"
    scp -P $ssh_port $ssh_opts $target:$VECONFDIR/$veid.conf /vz/private/$target_veid/ve.conf
    [ $? -ne 0 ] && error "Failed to copy Container config file"

    echo "Container $veid: Setting up destination Container..."

    # Modify VEID inside config
    if [ $target_veid != $veid ]; then
        printf "Destination: sed -e "s,^VEID=.*,VEID=\"$target_veid\",g" -i /vz/private/$target_veid/ve.conf\n"
        sed -e "s,^VEID=.*,VEID=\"$target_veid\",g" -i /vz/private/$target_veid/ve.conf
        [ $? -ne 0 ] && error "Failed to modify Container config file"
    fi

    # Fix VE_ROOT - proxmox case
    printf "Source %s: sed -e "s,^VE_ROOT=.*,VE_ROOT=\"/vz/root/$target_veid\",g" -i /vz/private/$target_veid/ve.conf\n", "$target"
    sed -e "s,^VE_ROOT=.*,VE_ROOT=\"/vz/root/$target_veid\",g" -i /vz/private/$target_veid/ve.conf
    [ $? -ne 0 ] && error "Failed to fix VE_ROOT in Container config file"

    # Adjust disk space on destination
    printf "Source %s: sed -e "s,^DISKSPACE=.*,DISKSPACE=\"$required_space:$required_space\",g" -i /vz/private/$target_veid/ve.conf\n", "$target"
    sed -e "s,^DISKSPACE=.*,DISKSPACE=\"$required_space:$required_space\",g" -i /vz/private/$target_veid/ve.conf
    [ $? -ne 0 ] && error "Failed to fix destination Container diskspace"

    # Check for ostemplate on target
    eval `ssh -p $ssh_port $ssh_opts $target grep "^OSTEMPLATE=" $VECONFDIR/$veid.conf`
    ostemplate_rpm=`ssh -p $ssh_port $ssh_opts $target echo $OSTEMPLATE | sed "s,^\.,,g"`-ez
    printf "Destination: rpm -q $ostemplate_rpm\n"
    rpm -q $ostemplate_rpm
    if [ $? -ne 0 ]; then
        printf "Destination: yum install -y $ostemplate_rpm\n"
        yum install -y $ostemplate_rpm
        if [ $? -ne 0 ]; then
            # And disable template if not supported
            printf "Destination: sed -e "s,^OSTEMPLATE=,#OSTEMPLATE=,g" -i /vz/private/$target_veid/ve.conf\n"
            sed -e "s,^OSTEMPLATE=,#OSTEMPLATE=,g" -i /vz/private/$target_veid/ve.conf
            [ $? -ne 0 ] && error "Failed to disable param OSTEMPLATE in Container config file"
        fi
    fi

    # Comment out QUOTATIME
    printf "Destination: sed -e "s,^QUOTATIME=,#QUOTATIME=,g" -i /vz/private/$target_veid/ve.conf\n"
    sed -e "s,^QUOTATIME=,#QUOTATIME=,g" -i /vz/private/$target_veid/ve.conf

    # Comment out QUOTAUGIDLIMIT
    printf "Destination: sed -e "s,^QUOTAUGIDLIMIT=,#QUOTAUGIDLIMIT=,g" -i /vz/private/$target_veid/ve.conf\n"
    sed -e "s,^QUOTAUGIDLIMIT=,#QUOTAUGIDLIMIT=,g" -i /vz/private/$target_veid/ve.conf

    # Stop source
    printf "Source %s: ssh -p $ssh_port $ssh_opts $target vzctl stop $veid", "$target"
    ssh -p $ssh_port $ssh_opts $target vzctl stop $veid

    # Remove bind_mount
    printf "Destination: ssh -p $ssh_port $ssh_opts $target vzctl exec $veid umount $tmpdir\n"
    ssh -p $ssh_port $ssh_opts $target vzctl exec $veid umount $tmpdir
    printf "Destination: ssh -p $ssh_port $ssh_opts $target vzctl exec $veid rm -rf $tmpdir\n"
    ssh -p $ssh_port $ssh_opts $target vzctl exec $veid rm -rf $tmpdir

    # Register Container on target
    printf "Destination: vzctl register /vz/private/$target_veid $target_veid\n"
    vzctl register /vz/private/$target_veid $target_veid
    [ $? -ne 0 ] && error "Failed to register Container $target_veid"

    # Mount target
    printf "Destination: vzctl mount $target_veid\n"
    vzctl mount $target_veid
    [ $? -ne 0 ] && error "Failed to mount destination Container"

    # Remove
    printf "Destination: rm -f /vz/root/$target_veid/aquota.*\n"
    rm -f /vz/root/$target_veid/aquota.*
    printf "Destination: find /vz/root/$target_veid/etc -name *vzquota* | xargs rm -f\n"
    find /vz/root/$target_veid/etc -name *vzquota* | xargs rm -f

    # Change
    printf "Destination: rm -f /vz/root/$target_veid/etc/mtab\n"
    rm -f /vz/root/$target_veid/etc/mtab
    printf "Destination: ln -s /proc/mounts /vz/root/$target_veid/etc/mtab\n"
    ln -s /proc/mounts /vz/root/$target_veid/etc/mtab
    [ $? -ne 0 ] && error "Failed to fix /etc/mtab in destination Container"

    # Create devices
    echo 'Creating devices'
    printf "Destination: mknod /vz/root/$target_veid/dev/ptmx c 5 2\n"
    mknod /vz/root/$target_veid/dev/ptmx c 5 2
    printf "Destination: chmod 666 /vz/root/$target_veid/dev/ptmx\n"
    chmod 666 /vz/root/$target_veid/dev/ptmx
    printf "Destination: mknod /vz/root/$target_veid/etc/udev/devices/ptmx c 5 2\n"
    mknod /vz/root/$target_veid/etc/udev/devices/ptmx c 5 2

    for i in 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15; do
        a=`awk -v num=$i 'BEGIN { printf("%x\n", num) }'`
        mknod /vz/root/$target_veid/dev/ttyp$a c 3 $i
        mknod /vz/root/$target_veid/etc/udev/devices/ttyp$a c 3 $i
        mknod /vz/root/$target_veid/dev/ptyp$a c 2 $i
        mknod /vz/root/$target_veid/etc/udev/devices/ptyp$a c 2 $i
    done

    # Try to restore quotas if any
    echo 'Trying to restore quotas'
    if [ "x$quota_restore_command" != "x" ]; then
        # Start target
        printf "Destination: vzctl start $target_veid\n"
        vzctl start $target_veid > /dev/null 2>&1
        [ $? -ne 0 ] && error "Failed to start destination Container"

        # Restore quotas
        printf "Destination: vzctl exec $target_veid "$quota_restore_command"\n"
        vzctl exec $target_veid "$quota_restore_command"
        [ $? -ne 0 ] && error "Failed to restore quota"

        # Stop target
        printf "Destination: vzctl stop $target_veid\n"
        vzctl stop $target_veid
    else
        printf "Destination: vzctl umount $target_veid\n"
        vzctl umount $target_veid
    fi

    echo "Container $veid: Done, cleaning..."
}

# Check for parameters
TARGET=$1
SSH_PORT=$2
SSH_OPTS=$3
shift 3
COUNT=0

while [ ! -z $1 ]; do
	VEID_LIST[$COUNT]=${1%:*}
	[ -z ${VEID_LIST[$COUNT]} ] && usage
	TARGET_VEID_LIST[$COUNT]=${1#*:}
	[ -z ${TARGET_VEID_LIST[$COUNT]} ] && TARGET_VEID_LIST[$COUNT]=${VEID_LIST[$COUNT]}
	shift
	COUNT=$((COUNT+1))
done

[ -z "$TARGET" -o ${#VEID_LIST[@]} -eq 0 ] && usage

echo "Checking target host parameters..."

# Check for ssh -p $ssh_port $ssh_opts key
ssh -p $SSH_PORT $SSH_OPTS -o PasswordAuthentication="no" $TARGET exit
[ $? -ne 0 ] && error "ssh -p $SSH_PORT $SSH_OPTS key was not configured for $TARGET"

MIGRATION_STARTED=1

# Get path to config dir
if [ -d /etc/sysconfig/vz-scripts ]; then
	VECONFDIR=/etc/sysconfig/vz-scripts
elif [ -d /etc/vz/conf ]; then
	VECONFDIR=/etc/vz/conf
else
	error "Failed to detect Containers config dir"
fi

# Start migration in parallel
while [ $COUNT -gt 0 ]; do
	COUNT=$((COUNT-1))
	(migrate ${VEID_LIST[$COUNT]} ${TARGET_VEID_LIST[$COUNT]} $TARGET $SSH_PORT "$SSH_OPTS" ;) &
	PIDS_LIST[$!]=${VEID_LIST[$COUNT]}
	echo "Migration of Container ${VEID_LIST[$COUNT]} started"
done

PIDS=`jobs -p`
RC=0

while [ ! -z "$PIDS" ]; do
	for pid in $PIDS; do
		sleep 1
		kill -0 $pid
		[ $? -eq 0 ] && continue
		wait $pid
		err=$?

		if [ $err -ne 0 ]; then
		    echo "migrationCheck:1"
			echo "Migraton of Container ${PIDS_LIST[$pid]} failed with error $err"
			RC=1
		else
		    echo "migrationCheck:0"
            echo "Migration of Container ${PIDS_LIST[$pid]} completed successfully"
		fi
		PIDS=`echo $PIDS | sed "s,$pid,,g"`
	done
done

echo "All done"

exit $RC
